This project is work in progress
equinox_vision
is a library for small-size image classification problems in JAX with equinox.
pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install -U -e .
See examples.