My Personal JAX Library

This is my personal library where I have all of my code that uses jax and the equinox backend.

Pronounced: Jay-EE-JEX

---

Components

Versioning: versioneer

Documentation: jupyterbook

---

Installation

This package isn't pip-worthy (yet) but here are a few options for installation.

Option I: Use the pip install option (locally)

https://github.com/jejjohnson/jemanjjax.git
cd jemanjjax
pip install -e .[dev,all]

Option II: Install it from pip directly.

pip install "git+https://github.com/jejjohnson/jemanjjax.git"

---

External Packages

I use quite a few of external packages that I've relegated to their own repo.

Neural Fields

pip install "git+https://github.com/jejjohnson/eqx-nerf.git"

Trainer

pip install "git+https://github.com/jejjohnson/eqx-trainer.git"

OceanBench

brew install g++ cmake eigen boost gsl
pip install "git+https://github.com/jejjohnson/oceanbench.git"

---

---

Inspiration

• UVADLC Course
• JaxLightning
• Tez