Metadata-Version: 2.4
Name: jaxpy
Version: 0.0.0
Summary: Add your description here
Author-email: Bernardo Esteves <24633851+esteveste@users.noreply.github.com>
Requires-Python: >=3.10
Requires-Dist: flax
Requires-Dist: jax
Description-Content-Type: text/markdown

<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
```

