PyTorch-like neural network module system in JAX.
The idea is borrowed from Parallax (https://github.com/srush/parallax). The implementation is different.
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