maxtext VS flax

This might be a tangent, but why does JAX only support the saving / serialization of AOT compilation executables for TPU [1]? It would be great to have the ability to save compiled functions and not have to JIT compile something every time you restart a session.

(Julia used to have this problem too, but they've made great progress on caching JIT compiled functions to reduce latency.)

[1]: https://github.com/google/maxtext?tab=readme-ov-file#ahead-o...

Is t5x an encoder/decoder architecture?

Some more general options.

The Flax ecosystem

https://github.com/google/flax?tab=readme-ov-file

or dm-haiku

https://github.com/google-deepmind/dm-haiku

were some of the best developed communities in the Jax AI field

Perhaps the “trax” repo? https://github.com/google/trax

Some HF examples https://github.com/huggingface/transformers/tree/main/exampl...

Sadly it seems much of the work is proprietary these days, but one example could be Grok-1, if you customize the details. https://github.com/xai-org/grok-1/blob/main/run.py

https://github.com/google/flax/discussions/919 https://flax.readthedocs.io/en/latest/_modules/flax/linen/attention.html

Just as an FYI, you might be competing against another big open source project with the same name https://github.com/google/flax

I have asked the same question on the Flax discussion page on Github as well.

What's your favorite Deep Learning API for JAX - Flax, Haiku, Elegy, something else?

As a researcher in RL & ML in a big industry lab, I would say most of my colleagues are moving to JAX 0https://github.com/google/jax], which this article kind of ignores. JAX is XLA-accelerated NumPy, it's cool beyond just machine learning, but only provides low-level linear algebra abstractions. However you can put something like Haiku [https://github.com/deepmind/dm-haiku] or Flax [https://github.com/google/flax] on top of it and get what the cool kids are using :)

- https://github.com/google/flax/tree/main/examples

Any thoughts on Flax?