Jax codebases

Enumerating some codebases that use jax (and/or flax) in an elegant way.

• @ikostrikov’s codebases.
• jaxrl_m
• jaxrl
• Dataset iterator
• Jax models
• Jax Meta