Python reference#

Core transforms#

State-aware JAX transformations backed by Flax NNX.

class jaxsaw.core.Rngs(*args: Any, **kwargs: Any)#: Named, checkpointable random streams, e.g. Rngs(params=0, dropout=1).

jaxsaw.core.apply(graph, state, *args, **kwargs)#: Pure application for deterministic modules; returns output and updated state.

jaxsaw.core.stop_gradient(tree)#: Detach every array in a pytree, including recurrent state tuples.

Module composition#

Composable modules with explicit parameters and non-trainable buffers.

class jaxsaw.nn.module.Module(*args: Any, **kwargs: Any)#

Subclass and implement forward; child modules are discovered by NNX.

eval()#

Sets the Module to evaluation mode.

eval uses set_attributes to recursively set attributes deterministic=True and use_running_average=True of all nested Modules that have these attributes. Its primarily used to control the runtime behavior of the Dropout and BatchNorm Modules.

Example:

>>> from flax import nnx
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     self.dropout = nnx.Dropout(0.5)
...     self.batch_norm = nnx.BatchNorm(10, rngs=rngs)
...
>>> block = Block(2, 5, rngs=nnx.Rngs(0))
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(False, False)
>>> block.eval()
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(True, True)

Parameters:: **attributes – additional attributes passed to set_attributes.

state_dict()#: Copy canonical parameter/buffer arrays to host; no executable objects.

train(mode=True)#

Sets the Module to training mode.

train uses set_attributes to recursively set attributes deterministic=False and use_running_average=False of all nested Modules that have these attributes. Its primarily used to control the runtime behavior of the Dropout and BatchNorm Modules.

Example:

>>> from flax import nnx
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     # initialize Dropout and BatchNorm in eval mode
...     self.dropout = nnx.Dropout(0.5, deterministic=True)
...     self.batch_norm = nnx.BatchNorm(10, use_running_average=True, rngs=rngs)
...
>>> block = Block(2, 5, rngs=nnx.Rngs(0))
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(True, True)
>>> block.train()
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(False, False)

Parameters:: **attributes – additional attributes passed to set_attributes.

class jaxsaw.nn.module.ModuleDict(*args: Any, **kwargs: Any)#

class jaxsaw.nn.module.ModuleList(*args: Any, **kwargs: Any)#: Registered ordered modules; iteration does not call them.

class jaxsaw.nn.module.ParameterSet(module)#: Live parameter selection owned by a module (not an array snapshot).

class jaxsaw.nn.module.Sequential(*args: Any, **kwargs: Any)#: Chain deterministic single-input layers. Use forward for RNG-aware branches.

Layers#

Dense, convolutional, activation, and normalization layers.

class jaxsaw.nn.layers.AdaptiveAvgPool2d(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.AvgPool2d(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.BatchNorm1d(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.BatchNorm2d(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.Conv1d(*args: Any, **kwargs: Any)#: NCL convolution implemented through the same checked Conv2d primitive.

class jaxsaw.nn.layers.Conv2d(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.ConvTranspose2d(*args: Any, **kwargs: Any)#: Transposed NCHW convolution, with explicit integer padding/output_padding.

class jaxsaw.nn.layers.Dropout(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.Embedding(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.Flatten(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.GELU(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.LayerNorm(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.Linear(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.MaxPool2d(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.ReLU(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.Sigmoid(*args: Any, **kwargs: Any)#

class jaxsaw.nn.layers.Tanh(*args: Any, **kwargs: Any)#

Recurrent layers#

Batch-first recurrent layers; carries are explicit and sequence loops use scan.

class jaxsaw.nn.recurrent.LSTM(*args: Any, **kwargs: Any)#

cell_type#: alias of LSTMCell

class jaxsaw.nn.recurrent.LSTMCell(*args: Any, **kwargs: Any)#

class jaxsaw.nn.recurrent.RNN(*args: Any, **kwargs: Any)#

cell_type#: alias of RNNCell

class jaxsaw.nn.recurrent.RNNCell(*args: Any, **kwargs: Any)#

Attention#

Batch-first attention. Boolean mask True means the key is allowed.

class jaxsaw.nn.attention.MultiheadAttention(*args: Any, **kwargs: Any)#

Loss functions#

Stable loss functions and stateless activations; classifier inputs are logits.

Optimizers#

Optimizers with explicit gradients and live module parameter ownership.

Reference models#

Small reference architectures. These are teaching models, not pretrained weights.

class jaxsaw.models.Autoencoder(*args: Any, **kwargs: Any)#: MNIST autoencoder with named reconstruction and latent outputs.

class jaxsaw.models.ConvAutoencoder(*args: Any, **kwargs: Any)#: 28x28 reconstruction with a spatial latent tensor [B, channels, 7, 7].

class jaxsaw.models.DigitCNN(*args: Any, **kwargs: Any)#

Four convolution/pooling stages for MNIST NCHW images.

Zero padding expands 28x28 to 32x32 before four reductions to 16, 8, 4, and 2 pixels per side. Earlier two-stage checkpoints are incompatible.

class jaxsaw.models.Discriminator(*args: Any, **kwargs: Any)#

class jaxsaw.models.FashionDiscriminator(*args: Any, **kwargs: Any)#: Convolutional discriminator; outputs logits without a sigmoid.

class jaxsaw.models.FashionGenerator(*args: Any, **kwargs: Any)#: Convolutional Fashion-MNIST generator, with explicit BatchNorm state.

class jaxsaw.models.Generator(*args: Any, **kwargs: Any)#: Fashion-MNIST generator; output range [-1, 1].

class jaxsaw.models.LSTMClassifier(*args: Any, **kwargs: Any)#

class jaxsaw.models.MLP(*args: Any, **kwargs: Any)#

class jaxsaw.models.PetCNN(*args: Any, **kwargs: Any)#: RGB image classifier with global pooling and running normalization.

class jaxsaw.models.SequenceRegressor(*args: Any, **kwargs: Any)#: Next-value forecasting with an RNN or LSTM and explicit returned carry.

class jaxsaw.models.TinyLanguageModel(*args: Any, **kwargs: Any)#

class jaxsaw.models.TransformerBlock(*args: Any, **kwargs: Any)#

Data and training#

Small single-process datasets and reproducible batch loading.

class jaxsaw.data.DataLoader(dataset, batch_size=32, shuffle=False, seed=0, drop_last=False, collate_fn=<function collate>)#: Deterministic epoch shuffle; state records the next sample cursor.

Small reusable update steps; all optimization is explicit.

class jaxsaw.training.ModelCheckpoint(path, monitor='val/loss', mode='min')#: Save the best validation metric; checkpoint state remains explicit.

class jaxsaw.training.Task(model, optimizer, rngs=None)#

Optional experiment wrapper; training_step owns all optimizer updates.

Override training_step/validation_step. Declare persistent graph objects in checkpoint_objects; the default contains model, optimizer, and optional RNGs.

class jaxsaw.training.Trainer(max_epochs, callbacks=(), run=None)#: Host-side epoch loop delegating optimization to the task step.

jaxsaw.training.gan_step(generator, discriminator, g_optimizer, d_optimizer, real, rngs, latent_size=64)#: One D update then one G update; inactive model buffers stay in eval mode.

Checkpoints and inspection#

Atomic, non-pickle NPZ checkpoints for preconstructed matching object graphs.

jaxsaw.checkpoint.restore(path, **objects)#: Validate the entire snapshot before updating any live participant.

jaxsaw.checkpoint.save(path, *, metadata=None, progress=None, data_state=None, **objects)#

Save named models/optimizers/RNGs together, preserving joint graph paths.

Restore requires the same named participants, architecture and optimizer config. Training modes are recorded separately. There is no executable deserialization.

JSON-safe architecture hierarchy and non-mutating output shape inspection.

jaxsaw.inspection.describe(model)#: Describe ownership/configuration and aliases, not a complete execution graph.

jaxsaw.inspection.output_spec(model, *inputs, **kwargs)#: Run an evaluation snapshot; return only output shapes/dtypes.

Local structured run records; plotting libraries are optional consumers.