I recently started migrating from Tensorflow to Jax , Flax. I’ve created a new repository to help people transition from TensorFlow to JAX and Flax. It features a simple Flax trainer that’s perfect for beginners wanting to get their learn training models with the Flax framework.It’s open source, and I’d love to see contributions from the community. Thank you.