synax 0.0.0

A neural network library for JAX

Synax

A magic-free neural network library for JAX.

No tracing. No transforms. Just purely-functional JAX.

Installation

pip install git+https://github.com/carlosgmartin/synax

Editable install:

pip install -e .

Example

from jax import numpy as jnp, random
import synax

# Create a module.
module = synax.MLP([2, 32, 3])
# Create a random key.
key = random.key(0)
# Initialize parameters.
w = module.init(key)
# Define an input.
x = jnp.ones(2)
# Compute the output.
y = module.apply(w, x)
# Print the output.
print(y)

Codebase quality control

Run the following after every change:

ruff check && ruff format && pyright && pytest

Documentation

Build documentation:

sphinx-build -M html docs docs/_build