JAX is really cool, but still somewhat immature. I would love to see it taking more ground and improving wrt e.g. integration with tensorboard and getting all t

JAX is really cool, but still somewhat immature. I would love to see it taking more ground and improving wrt e.g. integration with tensorboard and getting all the goodies we have in tensorflow. If you are looking for a higher level framework, I would recommend elegy [0] which is very close to the keras API. [0] https://github.com/poets-ai/elegy https://github.com/poets-ai/elegy