Neural networks

Layers, blocks, inference sessions, trainers, optimizers and schedules.

Adagrad	Adagrad optimizer, as torch.optim.Adagrad.
Adam	Adam optimizer, as torch.optim.Adam, over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object).
Adamax	Adamax optimizer, as torch.optim.Adamax.
AdamW	AdamW optimizer, as torch.optim.AdamW, over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object).
clip_grad_norm_	Clip gradients by their global L2 norm, as torch.nn.utils.clip_grad_norm_.
cross_entropy	Cross-entropy loss over class indices, as torch.nn.functional.cross_entropy.
Embedding	Embedding table with a reproducible gather and dense gradient, as torch.nn.Embedding.
embedding	Embedding lookup tables with a reproducible backward pass.
ExponentialLR	Exponential learning-rate schedule: multiply by gamma every step.
GRUClassifier	GRU sequence classifier: recurrent layers under a linear head.
GRURegressor	GRU sequence regressor: recurrent layers under a linear head.
layer_norm_backward	Gradients of layer_norm_forward.
layer_norm_forward	Layer normalization over the trailing dimensions, as torch.nn.functional.layer_norm.
LayerNorm	Layer normalization layer, as torch.nn.LayerNorm.
Lion	Lion optimizer (Chen et al., 2023, "Symbolic Discovery of Optimization Algorithms").
LSTMClassifier	LSTM sequence classifier: recurrent layers under a linear head.
LSTMRegressor	LSTM sequence regressor: recurrent layers under a linear head.
mamba	Mamba-1, Mamba-2 and Mamba-3 blocks in float32.
Mamba1Block	One Mamba-1 block: RMSNorm, selective-scan mixer and residual.
Mamba1BlockInference	Forward-only Mamba1Block on the CPU, from weights trained on a GPU.
Mamba1State	Recurrent state of a Mamba1Block, owned by the caller.
Mamba2Block	One Mamba-2 (SSD) block: RMSNorm, mixer and residual.
Mamba2BlockInference	Forward-only Mamba2Block on the CPU, from weights trained on a GPU.
Mamba2State	Recurrent state of a Mamba2Block, owned by the caller.
Mamba3Block	One Mamba-3 (SISO) block: RMSNorm, mixer and residual.
Mamba3BlockInference	Forward-only Mamba3Block on the CPU, from weights trained on a GPU.
Mamba3State	Recurrent state of a Mamba3Block, owned by the caller.
MLPClassifier	Multi-layer perceptron classifier, as scikit-learn's MLPClassifier.
MLPInference	Forward-only SmallMLPTrainer network (8 -> 16 ReLU -> 3, float32) on the CPU.
MLPRegressor	Multi-layer perceptron regressor, as scikit-learn's MLPRegressor.
MoEBlock	Sparse mixture-of-experts feed-forward block, as Mixtral's MixtralSparseMoeBlock.
NAdam	NAdam optimizer, as torch.optim.NAdam.
neural_network	Small neural-network trainers with complete, resumable state.
OneCycleLR	One-cycle learning-rate schedule, as torch's OneCycleLR.
RMSprop	RMSprop optimizer, as torch.optim.RMSprop.
RNNClassifier	Elman RNN sequence classifier: recurrent layers under a linear head.
RNNRegressor	Elman RNN sequence regressor: recurrent layers under a linear head.
SambaConfig	Shape of a SambaStack: vocabulary, width and the sequence of layer kinds.
SambaInference	Forward-only SambaStack language model on the CPU.
SambaStack	A Samba-style language model with its own training loop.
SGD	Stochastic gradient descent, as torch.optim.SGD, over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object).
SmallMLPTrainer	Trainer for a fixed 8-16-3 ReLU network with mean cross-entropy and AdamW.
StepLR	Step-decay learning-rate schedule: multiply by gamma every step_size steps.
training	Optimizers, schedules, losses and training primitives over float32 arrays.
transformer	Llama-style transformer decoder block.
TransformerBlock	One Llama-shaped transformer decoder block, float32.
TransformerBlockInference	Forward-only TransformerBlock on the CPU, from weights trained on a GPU.
TransformerState	KV cache of a TransformerBlock, owned by the caller.