API and semantics#

Modules#

Subclass nn.Module, call super().__init__(), define children in the constructor, and implement forward. Call modules as model(x) so composition stays consistent. A child can return arrays, tuples, or dictionaries of arrays.

nn.Parameter(array) registers a differentiable variable; nn.Buffer(array) registers a non-parameter state variable. Access arrays through .value. Assigning the same module to two attributes shares weights. Parameter traversal deduplicates by identity. ModuleList and ModuleDict express named collections; the underlying NNX backend also discovers ordinary lists and dictionaries.

named_parameters(), named_buffers(), and named_modules() produce canonical paths. parameters() returns a live selection for an optimizer, not copies. state_dict() returns independent NumPy snapshots. A strict load_state_dict checks paths, shape, and dtype before copying values. Structural edits require a new optimizer and compiled step. Do not move/cast parameters after creating an optimizer; optimizer moments would otherwise retain the old device/dtype.

train()/eval() recursively toggle stochastic/normalization behavior and return self. They do not turn autodiff on/off. Switch mode outside a compiled step.

Layers and shapes#

Layer	Signature and convention
`Linear(I, O, bias=True, rngs=...)`	`[..., I] -> [..., O]`, weight `[O,I]`
`Conv1d(I, O, K, stride=1, padding=0, groups=1, rngs=...)`	`[B,C,L]`; integer padding
`Conv2d(I, O, K, stride=1, padding=0, groups=1, rngs=...)`	NCHW; integer/pair or SAME/VALID padding
`ConvTranspose2d(I, O, K, stride=1, padding=0, output_padding=0, rngs=...)`	NCHW; grouping and dilation unsupported; weights `[O,I,KH,KW]`
`MaxPool2d(K, stride=None)`, `AvgPool2d(...)`	NCHW, valid windows
`AdaptiveAvgPool2d((1,1))`	Global spatial mean; other sizes rejected
`BatchNorm1d(C)`, `BatchNorm2d(C)`	NC/NCL or NCHW; affine parameters and running buffers
`LayerNorm(shape)`	Normalize trailing dimensions
`Dropout(p)`	Explicit `rngs` in training; identity in evaluation
`Embedding(V,D,rngs=...)`	Integer `[...] -> [...,D]`
`RNN(I,H,num_layers=1, rngs=...)`	`[B,T,I] -> (sequence, h)`
`LSTM(I,H,num_layers=1, rngs=...)`	`[B,T,I] -> (sequence, (h,c))`
`MultiheadAttention(D, heads, rngs=...)`	Query `[B,Tq,D]`, key/value `[B,Tk,D]`

ReLU, GELU, Tanh, Sigmoid, Flatten, and Sequential provide simple composition. Sequential forwards one value without RNG keywords; write a custom forward for stochastic layers or multiple inputs. Linear/conv use fan-in-scaled uniform weights and zero biases. LSTM forget bias is one. BatchNorm uses population batch variance for both normalization and running-statistic updates, not PyTorch’s unbiased running-variance update. Results can therefore differ from PyTorch.

Transposed convolution output length on each spatial axis is (input - 1) * stride - 2 * padding + kernel_size + output_padding.

Recurrent state#

Hidden states use [layers,B,H]. state=None initializes zeros. LSTM state is a pair (h,c). lengths[B] freezes the carry after valid timesteps and zeros padded outputs. Zero-length examples keep their initial carry. Lengths must be integers between zero and T; validate data before entering compiled code. Bidirectionality is not implemented. For streaming, pass the returned carry to the next chunk; for truncated backpropagation, call js.stop_gradient(carry) between chunks.

Attention#

Masks are boolean and broadcast to [B,heads,Tq,Tk]; True means allowed. nn.attention.causal_mask(T) has shape [1,1,T,T] and padding_mask(valid[B,T]) has shape [B,1,1,T]. Padded query positions still need a loss mask. Fully masked rows return zero weights and output, including projection bias, with finite gradients. need_weights=True returns (output, weights) and must be static under compilation. Width must be divisible by heads. Embedding indices and class targets must be in range; these are input-data preconditions, not runtime checks inserted into compiled kernels.

Randomness and compilation#

Constructors consume js.Rngs(params=seed). Runtime dropout consumes a separate js.Rngs(dropout=seed). rngs.next(name) advances a named checkpointable stream; missing names are errors. Sampling can use js.random (JAX’s random namespace).

js.jit, js.value_and_grad, and js.grad are graph-aware NNX transforms. Pass stateful objects explicitly. Use ordinary JAX transforms only on pure array functions or the split/merge/apply/update boundary. Static shapes/modes can recompile. Avoid arbitrary Python side effects inside a transformed function. value_and_grad differentiates parameter leaves of the selected module and returns an NNX gradient state; has_aux=True returns ((loss, aux), grads).

NNX semantics govern exceptions and advanced transform nesting. Eager mutations are not automatically rolled back after an exception. Keep shapes/configuration valid before calling transformed code.

Losses and optimization#

nn.functional exposes cross_entropy, binary_cross_entropy_with_logits, mse_loss, l1_loss, and smooth_l1_loss. Prediction/target shapes must match except cross-entropy’s class axis. Use integer labels with raw class logits; class axis defaults to last. Reductions are none, sum, mean. Masked cross-entropy divides by valid count, with zero loss/gradient for an all-masked batch. BCE uses logits for numerical stability.

optim.SGD, Adam, and AdamW accept model.parameters() and lr. optimizer.step(grads) updates once; optimizer.step_count reports the counter. clip_grad_norm(grads, max_norm) returns (clipped_grads, original_norm). There are no .grad buffers, implicit zeroing, parameter groups, or freezing API in this release. Use disjoint modules/optimizers for GANs and differentiate only the intended module. Gradients must originate from the same model structure; matching-shaped trees from an unrelated model are not identity-validated.

Training utilities#

training includes compiled steps for classification, CNN dropout, forecasting, autoencoding, and the two GAN phases. They are readable reference implementations, not a mandatory trainer abstraction. Use your own step for custom objectives.

Task(model, optimizer, rngs=None) may override training_step, validation_step, set_phase, and checkpoint_objects. Steps own updates; Trainer never differentiates again. Trainer(max_epochs, callbacks=(), run=None) aggregates loss, count, and optional correct; custom metrics/GAN orchestration should use the explicit loop. ModelCheckpoint(path, monitor='val/loss', mode='min') saves the selected checkpoint. Automatic trainer resume and callback-state resume are not implemented; use the explicit checkpoint/loader APIs.