{"library":"orbax-export","title":"Orbax Export","description":"Orbax Export is a Python library designed for JAX users, providing utilities to serialize and export JAX models into the TensorFlow SavedModel format. It is a specialized component within the broader Orbax ecosystem, which offers common training utilities for JAX. The library is actively maintained, with its latest version being 0.0.8, released in September 2025, and maintains a regular release cadence to support JAX users.","language":"python","status":"active","last_verified":"Fri May 15","install":{"commands":["pip install orbax-export","pip install orbax-export[all]"],"cli":null},"imports":["from orbax.export import ExportManager","from orbax.export import JaxModule","from orbax.export import ServingConfig","import orbax.export"],"auth":{"required":false,"env_vars":[]},"quickstart":{"code":"import os\nimport jax\nimport jax.numpy as jnp\nimport tensorflow as tf # Required for SavedModel export\nfrom orbax.export import ExportManager, JaxModule, ServingConfig\n\n# Dummy JAX model and parameters for demonstration\nclass SimpleJAXModel:\n    def apply(self, params, inputs):\n        return params['w'] * inputs + params['b']\n\nmodel_instance = SimpleJAXModel()\nfinal_model_params_to_save = {'w': jnp.array(2.0), 'b': jnp.array(1.0)}\n\n# JAX Apply Function: The core JAX logic for the model's forward pass.\n@jax.jit\ndef jax_model_apply_fn_for_export(params, inputs):\n    return model_instance.apply(params, inputs)\n\n# Optional: TF Pre-processing Function\ndef tf_preprocess_fn_for_export(input_tensor: tf.Tensor) -> tf.Tensor:\n    return tf.cast(input_tensor, tf.float32) / 255.0\n\n# Optional: TF Post-processing Function\ndef tf_postprocess_fn_for_export(output_tensor: tf.Tensor) -> dict[str, tf.Tensor]:\n    return {'output': output_tensor}\n\n# Create a JaxModule\njax_module = JaxModule(\n    apply_fn=jax_model_apply_fn_for_export,\n    params=final_model_params_to_save,\n    preprocess_fn=tf_preprocess_fn_for_export,\n    postprocess_fn=tf_postprocess_fn_for_export\n)\n\n# Define serving signatures\nserving_signatures = {\n    'serving_default': ServingConfig(\n        input_signature=[\n            tf.TensorSpec(shape=[None, 1], dtype=tf.int32, name='input')\n        ],\n        output_signature={\n            'output': tf.TensorSpec(shape=[None, 1], dtype=tf.float32, name='output')\n        }\n    )\n}\n\n# Define export path\nexport_path = os.environ.get('ORBAX_EXPORT_PATH', '/tmp/my_jax_model_export')\n\n# Export the model\nexport_manager = ExportManager(\n    jax_module,\n    serving_signatures=serving_signatures\n)\nexport_manager.save(export_path)\n\nprint(f\"JAX model exported to TensorFlow SavedModel at: {export_path}\")\n\n# Basic verification (optional)\nloaded_model = tf.saved_model.load(export_path)\ninput_data = tf.constant([[5]], dtype=tf.int32)\noutput = loaded_model.signatures['serving_default'](input_data)\nprint(f\"Loaded model output for input {input_data.numpy()}: {output['output'].numpy()}\")\n","lang":"python","description":"This quickstart demonstrates how to define a simple JAX model, wrap it with `JaxModule`, configure serving signatures with `ServingConfig`, and then export it to the TensorFlow SavedModel format using `ExportManager.save()`. It includes optional TensorFlow pre- and post-processing functions that will be integrated into the SavedModel graph. Ensure `tensorflow` is installed to run this example.","tag":null,"tag_description":null,"last_tested":null,"results":[]},"compatibility":{"tag":null,"tag_description":null,"last_tested":"2026-05-15","installed_version":"0.0.5","pypi_latest":"0.0.8","is_stale":true,"summary":{"python_range":"3.10–3.9","success_rate":100,"avg_install_s":18.9,"avg_import_s":13.92,"wheel_type":"wheel"},"results":[{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"17.8M"},{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"17.8M"},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.6,"import_time_s":null,"mem_mb":null,"disk_size":"683M"},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":17.5,"import_time_s":null,"mem_mb":null,"disk_size":"683M"},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"19.6M"},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"19.6M"},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.7,"import_time_s":null,"mem_mb":null,"disk_size":"738M"},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.8,"import_time_s":null,"mem_mb":null,"disk_size":"738M"},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"11.5M"},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"11.5M"},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.3,"import_time_s":null,"mem_mb":null,"disk_size":"725M"},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.1,"import_time_s":null,"mem_mb":null,"disk_size":"725M"},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"11.2M"},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"11.2M"},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.5,"import_time_s":null,"mem_mb":null,"disk_size":"724M"},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15,"import_time_s":null,"mem_mb":null,"disk_size":"724M"},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"17.3M"},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":"17.3M"},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"orbax-export","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"broken","install_time_s":15.7,"import_time_s":null,"mem_mb":null,"disk_size":"623M"},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"all","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"clean","install_time_s":47.7,"import_time_s":13.92,"mem_mb":152.5,"disk_size":"2.6G"}]}}