How to serve aJAX model

[supercoderhawk]

Is there a good practice to serve JAX model, such as serving of tensorflow ?

[cgarciae]

Checkout jax2tf. You can use it to create a SavedModel, then use all the TF tooling for serving.

[8bitmp3]

• 1 @cgarciae

Check out these JAX to TensorFlow docs here:

• https://github.com/google/jax/blob/main/jax/experimental/jax2tf/examples/README.md (by @gnecula)
• https://dm-haiku.readthedocs.io/en/latest/notebooks/jax2tf.html or https://github.com/deepmind/dm-haiku/blob/main/docs/notebooks/jax2tf.ipynb (by @tomhennigan)