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 abil

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 has had 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-of-time-compilation-aot-tpu-only https://github.com/google/maxtext?tab=readme-ov-file#ahead-o...