nerva_jax documentation

A tiny, educational set of neural network components built on JAX.

Install and build

# from repository root
python -m pip install -U sphinx sphinx-rtd-theme
# build HTML docs into docs_sphinx/_build/html
sphinx-build -b html docs_sphinx docs_sphinx/_build/html

API reference

nerva_jax	nerva_jax: Minimal neural network components built on top of JAX.
nerva_jax.activation_functions	Activation functions and utilities used by the MLP implementation.
nerva_jax.datasets	In-memory data loader helpers and one-hot conversions.
nerva_jax.layers	Neural network layers used by the MultilayerPerceptron class.
nerva_jax.learning_rate	Learning-rate schedulers.
nerva_jax.loss_functions	Analytic loss functions and their gradients used during training.
nerva_jax.matrix_operations	Matrix operations built on top of torch to support the math in the library.
nerva_jax.multilayer_perceptron	A simple multilayer perceptron (MLP) class.
nerva_jax.optimizers	Optimizers used to adjusts the model's parameters based on the gradients.
nerva_jax.softmax_functions	Softmax and log-softmax functions together with stable variants.
nerva_jax.training	Training helpers for the MLP, including a basic SGD loop and CLI glue.
nerva_jax.utilities	Miscellaneous utilities (formatting, timing, parsing, I/O).
nerva_jax.weight_initializers	Weight and bias initialization helpers for linear layers.