<div align="center">
<img src="images/jaxpy.png" width="400" alt="logo"></img>
</div>

-----------------

jaxpy is a python package that:
+ makes jax easier to use with full functionality (from a pytorch user point of view)
+ automatic reproducible global random generator
+ integrated with the statefull class implementation of flax nnx


It integrates jax, flax and optax all in one package.

### Install
```bash
pip install jaxpy
```

