Yeah inference is a real weakness. I've used jax2tf to get tf saved models and tflite models, but it was a real journey... There's some effort going into syste

Yeah inference is a real weakness. I've used jax2tf to get tf saved models and tflite models, but it was a real journey... There's some effort going into systems for saving and restoring the computation graph for Jax programs, which will help a lot. I'm surprised it didn't happen sooner, as it seems like quite a natural fit with the jax architecture.