Embedded inference / ONNX?

[bhchiang]

How portable are checkpointed models - is it possible to convert to ONNX?

Use case: doing real-time inference with a trained Linen model. I realize that JAX is primarily geared towards research use cases so this may fall outside of the primary scope of Flax as well.

Thanks!

[avital]

Hi @bryanhpchiang -- I'm not sure if folks have tried converting JAX models to ONNX, but I can ask. The current best practice for inference is to use jax2tf and existing infrastructure for tf.SavedModels. (We should probably add a short guide about that on our documentation)

[bhchiang]

Thanks for suggesting jax2tf - TF conversion looks promising and would be a strong starting point. Short guide would be great!