Note that, in this space, there is also Flax[0] which is also built on top of Jax bringing more deep-learning specific primitives (while not trying to be tensor

Note that, in this space, there is also Flax[0] which is also built on top of Jax bringing more deep-learning specific primitives (while not trying to be tensorflow compatible unlike Trax if I understand correctly). [0]: https://github.com/google-research/flax/tree/prerelease https://github.com/google-research/flax/tree/prerelease