Need to move a JAX model into production infrastructure? JAX2TF converts a JAX model into a TensorFlow SavedModel, connecting JAX development with TensorFlow’s deployment ecosystem. section 2.8 section 12.1.3
A converted model can support several workflows:
- Serve inference on a server with TF Serving, on a device with TFLite, or on the web with TensorFlow.js. section 2.8 section 12.1.3
- Continue fine-tuning a JAX-trained model in TensorFlow with existing training data and setup. section 2.8
- Combine components trained with JAX and TensorFlow in a fused model. section 12.1.3
From JAX training to serving
Train a model in JAX, convert it to a TensorFlow SavedModel with JAX2TF, and deploy inference through TF Serving, TFLite, or TensorFlow.js. Remember that a saved checkpoint may contain only model weights, so the prediction function and its code still need to be reproduced or incorporated into the deployment workflow. section 2.8 section 12.1.3
