pax3 (pypi): A stateful pytree library for training neural networks.
Keywords: deep-learning|jax|opax|pytree