Nmax

PyTorch-like neural network module system in JAX.

The idea is borrowed from Parallax (https://github.com/srush/parallax). The implementation is different.

Instalation

git clone https://github.com/toru34/nmax
cd nmax
python setup.py sdist bdist_wheel
pip install ./dist/nmax-0.0.1.tar.gz

See examples/xor/xor.py