Miscellaneous extensions for JAX that might be too niche to upstream.
Note: This is just a dumping ground for utilities I find myself re-using, so the API is subject to change and I make no guarantees about the stability of the code!
pip install git+https://github.com/adrn/jax-ext