xax 0.2.20__tar.gz → 0.2.22__tar.gz

{xax-0.2.20/xax.egg-info → xax-0.2.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.20
+Version: 0.2.22
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,22 +31,6 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
-Provides-Extra: exportable
-Requires-Dist: flax; extra == "exportable"
-Requires-Dist: orbax-export; extra == "exportable"
-Requires-Dist: tensorflow; extra == "exportable"
-Provides-Extra: all
-Requires-Dist: black; extra == "all"
-Requires-Dist: darglint; extra == "all"
-Requires-Dist: mypy; extra == "all"
-Requires-Dist: ruff; extra == "all"
-Requires-Dist: pytest; extra == "all"
-Requires-Dist: types-pillow; extra == "all"
-Requires-Dist: types-psutil; extra == "all"
-Requires-Dist: types-requests; extra == "all"
-Requires-Dist: flax; extra == "all"
-Requires-Dist: orbax-export; extra == "all"
-Requires-Dist: tensorflow; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.20 → xax-0.2.22}/pyproject.toml

@@ -39,7 +39,6 @@ module = [
     "setuptools.*",
     "tensorboard.*",
     "transformers.*",
-    "orbax.export.*",
 ]
 
 ignore_missing_imports = true

{xax-0.2.20 → xax-0.2.22}/setup.py

@@ -14,12 +14,6 @@ with open("xax/requirements.txt", "r", encoding="utf-8") as f:
 with open("xax/requirements-dev.txt", "r", encoding="utf-8") as f:
     requirements_dev: list[str] = f.read().splitlines()
 
-requirements_export: list[str] = [
-    "flax",
-    "orbax-export",
-    "tensorflow",
-]
-
 with open("xax/__init__.py", "r", encoding="utf-8") as fh:
     version_re = re.search(r"^__version__ = \"([^\"]*)\"", fh.read(), re.MULTILINE)
 assert version_re is not None, "Could not find version in xax/__init__.py"
@@ -39,8 +33,6 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
-        "exportable": requirements_export,
-        "all": requirements_dev + requirements_export,
     },
     package_data={
         "xax": [

{xax-0.2.20 → xax-0.2.22}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.20"
+__version__ = "0.2.22"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -34,12 +34,6 @@ __all__ = [
     "get_positional_embeddings",
     "get_rotary_embeddings",
     "rotary_embeddings",
-    "MLPHyperParams",
-    "export_eqx_mlp",
-    "load_eqx",
-    "load_eqx_mlp",
-    "make_eqx_mlp",
-    "save_eqx",
     "cubic_bezier_interpolation",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
@@ -118,8 +112,10 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
+    "grad",
     "jit",
     "scan",
+    "vmap",
     "save_jaxpr_dot",
     "ColoredFormatter",
     "configure_logging",
@@ -215,12 +211,6 @@ NAME_MAP: dict[str, str] = {
     "get_positional_embeddings": "nn.embeddings",
     "get_rotary_embeddings": "nn.embeddings",
     "rotary_embeddings": "nn.embeddings",
-    "MLPHyperParams": "nn.equinox",
-    "export_eqx_mlp": "nn.equinox",
-    "load_eqx": "nn.equinox",
-    "load_eqx_mlp": "nn.equinox",
-    "make_eqx_mlp": "nn.equinox",
-    "save_eqx": "nn.equinox",
     "cubic_bezier_interpolation": "nn.geom",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
@@ -299,8 +289,10 @@ NAME_MAP: dict[str, str] = {
     "save_config": "utils.experiments",
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
+    "grad": "utils.jax",
     "jit": "utils.jax",
     "scan": "utils.jax",
+    "vmap": "utils.jax",
     "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
@@ -392,16 +384,6 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_rotary_embeddings,
         rotary_embeddings,
     )
-    from xax.nn.equinox import (
-        DTYPE,
-        ActivationFunction,
-        MLPHyperParams,
-        export_eqx_mlp,
-        load_eqx,
-        load_eqx_mlp,
-        make_eqx_mlp,
-        save_eqx,
-    )
     from xax.nn.geom import (
         cubic_bezier_interpolation,
         euler_to_quat,
@@ -482,7 +464,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
-    from xax.utils.jax import jit, scan
+    from xax.utils.jax import grad, jit, scan, vmap
     from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,

{xax-0.2.20 → xax-0.2.22}/xax/cli/edit_config.py

@@ -52,14 +52,24 @@ def main() -> None:
             print(colored(line, "light-cyan"), flush=True)
 
     # Saves the edited config to the checkpoint.
-    with tarfile.open(args.ckpt_path, "w:gz") as tar:
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        with tarfile.open(args.ckpt_path, "r:gz") as src_tar:
+            for member in src_tar.getmembers():
+                if member.name != "config":  # Skip the old config file
+                    src_tar.extract(member, tmp_dir)
 
-        def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
-            info = tarfile.TarInfo(name=name)
-            info.size = len(data)
-            tar.addfile(info, io.BytesIO(data))
+        with tarfile.open(args.ckpt_path, "w:gz") as tar:
+            for root, _, files in os.walk(tmp_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, tmp_dir)
+                    tar.add(file_path, arcname=arcname)
 
-        add_file_bytes("config", edited_config_str.encode())
+            # Add the new config file
+            info = tarfile.TarInfo(name="config")
+            config_bytes = edited_config_str.encode()
+            info.size = len(config_bytes)
+            tar.addfile(info, io.BytesIO(config_bytes))
 
 
 if __name__ == "__main__":

{xax-0.2.20 → xax-0.2.22}/xax/nn/metrics.py

@@ -7,8 +7,6 @@ import jax
 import jax.numpy as jnp
 from jaxtyping import Array
 
-from xax.utils.jax import jit as xax_jit
-
 NormType = Literal["l1", "l2"]
 
 
@@ -36,7 +34,6 @@ def dynamic_time_warping(distance_matrix_nm: Array) -> Array: ...
 def dynamic_time_warping(distance_matrix_nm: Array, return_path: Literal[True]) -> tuple[Array, Array]: ...
 
 
-@xax_jit(static_argnames=["return_path"])
 def dynamic_time_warping(distance_matrix_nm: Array, return_path: bool = False) -> Array | tuple[Array, Array]:
     """Dynamic Time Warping.
 

{xax-0.2.20 → xax-0.2.22}/xax/task/mixins/train.py

@@ -625,9 +625,13 @@ class TrainMixin(
         grad_metrics = {"grad_norm": grad_norm}
 
         def apply(grads: PyTree, grad_norm: Array) -> tuple[PyTree, optax.OptState]:
-            # Clip the global gradient norm to some desired range.
-            grad_factor = self.config.global_grad_clip / jnp.maximum(grad_norm, 1e-6)
-            grads = jax.tree.map(lambda x: x * grad_factor, grads)
+            # Clip gradients based on global norm, similar to optax.clip_by_global_norm
+            trigger = jnp.squeeze(grad_norm < self.config.global_grad_clip)
+
+            def clip_fn(t: Array) -> Array:
+                return jax.lax.select(trigger, t, (t / grad_norm.astype(t.dtype)) * self.config.global_grad_clip)
+
+            grads = jax.tree.map(clip_fn, grads)
 
             # Apply the gradient updates.
             updates, new_opt_state = optimizer.update(grads, opt_state, model_arr)

{xax-0.2.20 → xax-0.2.22}/xax/utils/jax.py

@@ -6,13 +6,14 @@ import logging
 import os
 import time
 from functools import wraps
-from typing import Any, Callable, Iterable, ParamSpec, Sequence, TypeVar, cast
+from typing import Any, Callable, Hashable, Iterable, ParamSpec, Sequence, TypeVar, cast
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 from jax._src import sharding_impls
 from jax._src.lib import xla_client as xc
+from jaxtyping import PyTree
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +21,7 @@ DEFAULT_COMPILE_TIMEOUT = 1.0
 
 Number = int | float | np.ndarray | jnp.ndarray
 
+T = TypeVar("T", bound=PyTree)
 
 P = ParamSpec("P")  # For function parameters
 R = TypeVar("R")  # For function return type
@@ -29,6 +31,9 @@ Carry = TypeVar("Carry")
 X = TypeVar("X")
 Y = TypeVar("Y")
 
+F = TypeVar("F", bound=Callable)
+AxisName = Hashable
+
 
 @functools.lru_cache(maxsize=None)
 def disable_jit_level() -> int:
@@ -166,6 +171,22 @@ def jit(
     return decorator
 
 
+def _split_module(tree: T, axis: int = 0) -> list[T]:
+    """Splits a module in the same way that jax.lax.scan and jax.vmap do.
+
+    Args:
+        tree: The tree to split.
+        axis: The axis to split on.
+
+    Returns:
+        A list of the split trees.
+    """
+    first_leaf = jax.tree.leaves(tree)[0]
+    num_slices = first_leaf.shape[axis]
+    result = [jax.tree.map(lambda x, idx=i: jnp.take(x, idx, axis=axis), tree) for i in range(num_slices)]
+    return result
+
+
 def scan(
     f: Callable[[Carry, X], tuple[Carry, Y]],
     init: Carry,
@@ -195,15 +216,96 @@ def scan(
     if not should_disable_jit(jit_level):
         return jax.lax.scan(f, init, xs, length, reverse, unroll)
 
+    carry = init
+    ys = []
+
     if xs is None:
         if length is None:
             raise ValueError("length must be provided if xs is None")
-        xs = cast(X, [None] * length)
+        for _ in range(length) if not reverse else range(length - 1, -1, -1):
+            carry, y = f(carry, None)  # type: ignore[arg-type]
+            ys.append(y)
 
-    carry = init
-    ys = []
-    for x in cast(Iterable, xs):
-        carry, y = f(carry, x)
-        ys.append(y)
+    else:
+        xlist = _split_module(xs, axis=0)
+        if reverse:
+            xlist = xlist[::-1]
+        for x in xlist:
+            carry, y = f(carry, x)
+            ys.append(y)
+
+    if reverse:
+        ys = ys[::-1]
+
+    if not ys:
+        return carry, jnp.array([])  # type: ignore[return-value]
 
     return carry, jax.tree.map(lambda *ys: jnp.stack(ys), *ys)
+
+
+def vmap(
+    fun: Callable[P, R],
+    in_axes: int | Sequence[int | None] = 0,
+    jit_level: int | None = None,
+) -> Callable[P, R]:
+    """A wrapper around jax.lax.vmap that allows for more flexible tracing.
+
+    If the provided JIT level is below the environment JIT level, we manually
+    unroll the scan function as a for loop.
+    """
+    if not should_disable_jit(jit_level):
+        return jax.vmap(fun, in_axes=in_axes)
+
+    @functools.wraps(fun)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+        if kwargs:
+            raise ValueError("vmap does not support keyword arguments")
+
+        ia = in_axes
+        if isinstance(ia, int):
+            ia = [ia] * len(args)
+        elif len(ia) != len(args):
+            raise ValueError("in_axes must be the same length as args")
+
+        if not all(isinstance(a, int) or a is None for a in ia):
+            raise ValueError("in_axes must be a list of integers or None")
+
+        ns = next((len(_split_module(a, axis=i)) for i, a in zip(ia, args, strict=True) if i is not None), None)
+        if ns is None:
+            return fun(*args, **kwargs)
+        split_args = [[a] * ns if i is None else _split_module(a, axis=i) for i, a in zip(ia, args, strict=True)]
+        split_outputs = [fun(*sargs, **kwargs) for sargs in zip(*split_args, strict=True)]
+
+        if not split_outputs:
+            return jnp.array([])  # type: ignore[return-value]
+
+        return jax.tree.map(lambda *ys: jnp.stack(ys), *split_outputs)
+
+    return wrapped
+
+
+def grad(
+    fun: Callable[P, R],
+    argnums: int | Sequence[int] = 0,
+    has_aux: bool = False,
+    holomorphic: bool = False,
+    allow_int: bool = False,
+    reduce_axes: Sequence[AxisName] = (),
+    jit_level: int | None = None,
+) -> Callable:
+    """A wrapper around jax.grad that allows for more flexible tracing.
+
+    We don't do anything special here, we just manually evaluate the function
+    if the JIT level is below the environment JIT level.
+    """
+    if not should_disable_jit(jit_level):
+        return jax.grad(fun, argnums, has_aux, holomorphic, allow_int, reduce_axes)
+
+    @functools.wraps(fun)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> Callable:
+        # Evaluate the function once, then just return the gradient.
+        fun(*args, **kwargs)
+
+        return jax.grad(fun, argnums, has_aux, holomorphic, allow_int, reduce_axes)(*args, **kwargs)
+
+    return wrapped

{xax-0.2.20 → xax-0.2.22/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.20
+Version: 0.2.22
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,22 +31,6 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
-Provides-Extra: exportable
-Requires-Dist: flax; extra == "exportable"
-Requires-Dist: orbax-export; extra == "exportable"
-Requires-Dist: tensorflow; extra == "exportable"
-Provides-Extra: all
-Requires-Dist: black; extra == "all"
-Requires-Dist: darglint; extra == "all"
-Requires-Dist: mypy; extra == "all"
-Requires-Dist: ruff; extra == "all"
-Requires-Dist: pytest; extra == "all"
-Requires-Dist: types-pillow; extra == "all"
-Requires-Dist: types-psutil; extra == "all"
-Requires-Dist: types-requests; extra == "all"
-Requires-Dist: flax; extra == "all"
-Requires-Dist: orbax-export; extra == "all"
-Requires-Dist: tensorflow; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.20 → xax-0.2.22}/xax.egg-info/SOURCES.txt

@@ -21,8 +21,6 @@ xax/core/conf.py
 xax/core/state.py
 xax/nn/__init__.py
 xax/nn/embeddings.py
-xax/nn/equinox.py
-xax/nn/export.py
 xax/nn/functions.py
 xax/nn/geom.py
 xax/nn/losses.py

{xax-0.2.20 → xax-0.2.22}/xax.egg-info/requires.txt

@@ -14,19 +14,6 @@ tensorboard
 psutil
 requests
 
-[all]
-black
-darglint
-mypy
-ruff
-pytest
-types-pillow
-types-psutil
-types-requests
-flax
-orbax-export
-tensorflow
-
 [dev]
 black
 darglint
@@ -36,8 +23,3 @@ pytest
 types-pillow
 types-psutil
 types-requests
-
-[exportable]
-flax
-orbax-export
-tensorflow

xax-0.2.20/xax/nn/equinox.py

@@ -1,183 +0,0 @@
-"""Equinox utilities."""
-
-import json
-import logging
-from pathlib import Path
-from typing import Callable, Literal, TypedDict, cast
-
-import equinox as eqx
-import jax
-from jaxtyping import PRNGKeyArray
-
-logger = logging.getLogger(__name__)
-
-ActivationFunction = Literal[
-    "relu",
-    "tanh",
-    "celu",
-    "elu",
-    "gelu",
-    "glu",
-    "hard_sigmoid",
-    "hard_silu",
-    "hard_swish",
-    "hard_tanh",
-    "leaky_relu",
-    "log_sigmoid",
-    "log_softmax",
-    "logsumexp",
-    "relu6",
-    "selu",
-    "sigmoid",
-    "soft_sign",
-    "softmax",
-    "softplus",
-    "sparse_plus",
-    "sparse_sigmoid",
-    "silu",
-    "swish",
-    "squareplus",
-    "mish",
-    "identity",
-]
-
-DTYPE = Literal["float32", "float64"]
-
-DTYPE_MAP: dict[DTYPE, jax.numpy.dtype] = {
-    "float32": jax.numpy.float32,
-    "float64": jax.numpy.float64,
-}
-
-
-class MLPHyperParams(TypedDict):
-    """Hyperparameters of an Equinox MLP."""
-
-    in_size: int | Literal["scalar"]
-    out_size: int | Literal["scalar"]
-    width_size: int
-    depth: int
-    activation: ActivationFunction
-    final_activation: ActivationFunction
-    use_bias: bool
-    use_final_bias: bool
-    dtype: DTYPE
-
-
-def _infer_activation(activation: ActivationFunction) -> Callable:
-    if activation == "identity":
-        return lambda x: x
-    try:
-        return getattr(jax.nn, activation)
-    except AttributeError as err:
-        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`") from err
-
-
-def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> eqx.nn.MLP:
-    """Create an Equinox MLP from a set of hyperparameters.
-
-    Args:
-        hyperparams: The hyperparameters of the MLP.
-        key: The PRNG key to use for the MLP.
-    """
-    activation = _infer_activation(hyperparams["activation"])
-    final_activation = _infer_activation(hyperparams["final_activation"])
-    dtype = DTYPE_MAP[hyperparams["dtype"]]
-
-    return eqx.nn.MLP(
-        in_size=hyperparams["in_size"],
-        out_size=hyperparams["out_size"],
-        width_size=hyperparams["width_size"],
-        depth=hyperparams["depth"],
-        activation=activation,
-        final_activation=final_activation,
-        use_bias=hyperparams["use_bias"],
-        use_final_bias=hyperparams["use_final_bias"],
-        dtype=dtype,
-        key=key,
-    )
-
-
-def export_eqx_mlp(
-    model: eqx.nn.MLP,
-    output_path: str | Path,
-    dtype: jax.numpy.dtype | None = None,
-) -> None:
-    """Serialize an Equinox MLP to a .eqx file.
-
-    Args:
-        model: The JAX MLP to export.
-        output_path: The path to save the exported model.
-        dtype: The dtype of the model.
-    """
-    if dtype is None:
-        dtype = eqx._misc.default_floating_dtype()
-
-    activation = model.activation.__name__
-    final_activation = model.final_activation.__name__
-
-    if final_activation == "<lambda>":
-        logger.warning("Final activation is a lambda function. Assuming identity.")
-        final_activation = "identity"
-
-    # cast strings to ActivationFunction for type checking
-    activation = cast(ActivationFunction, activation)
-    final_activation = cast(ActivationFunction, final_activation)
-
-    if dtype not in DTYPE_MAP.values():
-        raise ValueError(f"Invalid dtype: {dtype}. Must be one of {DTYPE_MAP.values()}")
-
-    dtype = {v: k for k, v in DTYPE_MAP.items()}[dtype]
-
-    hyperparams: MLPHyperParams = {
-        "in_size": model.in_size,
-        "out_size": model.out_size,
-        "width_size": model.width_size,
-        "depth": model.depth,
-        "activation": activation,
-        "final_activation": final_activation,
-        "use_bias": model.use_bias,
-        "use_final_bias": model.use_final_bias,
-        "dtype": dtype,
-    }
-
-    with open(output_path, "wb") as f:
-        hyperparam_str = json.dumps(hyperparams)
-        f.write((hyperparam_str + "\n").encode(encoding="utf-8"))
-        eqx.tree_serialise_leaves(f, model)
-
-
-def save_eqx(
-    model: eqx.Module,
-    output_path: str | Path,
-) -> None:
-    """Serialize an Equinox module to a .eqx file.
-
-    Args:
-        model: The Equinox module to export.
-        output_path: The path to save the exported model.
-    """
-    with open(output_path, "wb") as f:
-        eqx.tree_serialise_leaves(f, model)
-
-
-def load_eqx(
-    model: eqx.Module,
-    eqx_file: str | Path,
-) -> eqx.Module:
-    """Deserialize an Equinox module from a .eqx file.
-
-    Args:
-        model: The Equinox module to load into.
-        eqx_file: The path to the .eqx file to load.
-    """
-    with open(eqx_file, "rb") as f:
-        return eqx.tree_deserialise_leaves(f, model)
-
-
-def load_eqx_mlp(
-    eqx_file: str | Path,
-) -> eqx.nn.MLP:
-    with open(eqx_file, "rb") as f:
-        hyperparams = json.loads(f.readline().decode(encoding="utf-8"))
-        model = make_eqx_mlp(hyperparams=hyperparams, key=jax.random.PRNGKey(0))
-        return eqx.tree_deserialise_leaves(f, model)

xax-0.2.20/xax/nn/export.py

@@ -1,154 +0,0 @@
-"""Export JAX functions to TensorFlow SavedModel format."""
-
-import logging
-from pathlib import Path
-from typing import Callable
-
-import jax
-from jaxtyping import Array, PyTree
-
-try:
-    import flax
-    import tensorflow as tf
-    from jax.experimental import jax2tf
-    from orbax.export import ExportManager, JaxModule, ServingConfig
-except ImportError as e:
-    raise ImportError(
-        "In order to export models, please install Xax with exportable dependencies, "
-        "using 'xax[exportable]` to install the required dependencies."
-    ) from e
-
-logger = logging.getLogger(__name__)
-
-
-def _run_infer(tf_module: tf.Module, input_shapes: list[tuple[int, ...]], batch_size: int | None) -> tf.Tensor:
-    """Warm up the model by running it once."""
-    if batch_size is not None:
-        test_inputs = [
-            jax.random.normal(jax.random.PRNGKey(42), (batch_size, *input_shape)) for input_shape in input_shapes
-        ]
-    else:
-        test_inputs = [jax.random.normal(jax.random.PRNGKey(42), (1, *input_shape)) for input_shape in input_shapes]
-    if not hasattr(tf_module, "infer"):
-        raise ValueError("Model does not have an infer method")
-    return tf_module.infer(*test_inputs)
-
-
-def export(
-    model: Callable,
-    input_shapes: list[tuple[int, ...]],
-    output_dir: str | Path = "export",
-    batch_size: int | None = None,
-) -> None:
-    """Export a JAX function to TensorFlow SavedModel.
-
-    Note: Tensorflow GraphDef can't be larger than 2GB - https://github.com/tensorflow/tensorflow/issues/51870
-    You can avoid this by saving model parameters as non-constants.
-
-    Args:
-        model: The JAX function to export.
-        input_shapes: The shape of the input tensors, excluding batch dimension.
-        output_dir: Directory to save the exported model.
-        batch_size: Optional batch dimension. If None, a polymorphic batch dimension is used.
-    """
-    tf_module = tf.Module()
-    # Create a polymorphic shape specification for each input
-    poly_spec = "(b, ...)" if batch_size is not None else "(None, ...)"
-    polymorphic_shapes = [poly_spec] * len(input_shapes)
-    tf_module.infer = tf.function(  # type: ignore [attr-defined]
-        jax2tf.convert(
-            model,
-            polymorphic_shapes=polymorphic_shapes,
-            # setting this to False will allow the model to run on platforms other than the one that exports the model
-            # https://github.com/jax-ml/jax/blob/051687dc4c899df3d95c30b812ade401d8b31166/jax/experimental/jax2tf/README.md?plain=1#L1342
-            # generally though I think native_serialization is recommended
-            native_serialization=False,
-            with_gradient=False,
-        ),
-        autograph=False,
-        input_signature=[tf.TensorSpec([batch_size] + list(input_shape), tf.float32) for input_shape in input_shapes],
-    )
-
-    # warm up the model
-    _run_infer(tf_module, input_shapes, batch_size)
-
-    logger.info("Exporting SavedModel to %s", output_dir)
-    tf.saved_model.save(
-        tf_module,
-        output_dir,
-    )
-
-
-def export_with_params(
-    model: Callable,
-    params: PyTree,
-    input_shapes: list[tuple[int, ...]],
-    output_dir: str | Path = "export",
-    batch_dim: int | None = None,
-) -> None:
-    """Export a JAX function that takes parameters to TensorFlow SavedModel.
-
-    Args:
-        model: The JAX function to export. Should take parameters as first argument.
-        params: The parameters to use for the model.
-        input_shapes: The shape of the input tensors, excluding batch dimension.
-        output_dir: Directory to save the exported model.
-        batch_dim: Optional batch dimension. If None, a polymorphic batch dimension is used.
-    """
-    param_vars = tf.nest.map_structure(tf.Variable, params)
-
-    converted_model = jax2tf.convert(model)
-
-    def model_fn(*inputs: PyTree) -> Array:
-        return converted_model(param_vars, *inputs)
-
-    tf_module = tf.Module()
-    tf_module._variables = tf.nest.flatten(param_vars)  # type: ignore [attr-defined]
-    tf_module.infer = tf.function(  # type: ignore [attr-defined]
-        model_fn,
-        jit_compile=True,
-        autograph=False,
-        input_signature=[tf.TensorSpec([batch_dim] + list(input_shape), tf.float32) for input_shape in input_shapes],
-    )
-
-    # warm up the model
-    _run_infer(tf_module, input_shapes, batch_dim)
-
-    logger.info("Exporting SavedModel to %s", output_dir)
-    tf.saved_model.save(tf_module, output_dir)
-
-
-def export_flax(
-    model: flax.linen.Module,
-    params: PyTree,
-    input_shape: tuple[int, ...],
-    preprocessor: Callable | None = None,
-    postprocessor: Callable | None = None,
-    input_name: str = "inputs",
-    output_name: str = "outputs",
-    output_dir: str | Path = "export",
-) -> None:
-    jax_module = JaxModule(
-        params, model.apply, trainable=False, input_polymorphic_shape="(b, ...)"
-    )  # if you want to use a batch dimension
-
-    # to avoid mapping sequences to ambiguous mappings
-    if postprocessor is None:
-
-        def postprocessor(x: PyTree) -> PyTree:
-            return {output_name: x}
-
-    export_manager = ExportManager(
-        jax_module,
-        [
-            ServingConfig(
-                "serving_default",
-                input_signature=[tf.TensorSpec([None] + list(input_shape), tf.float32, name=input_name)],
-                tf_preprocessor=preprocessor,
-                tf_postprocessor=postprocessor,
-            )
-        ],
-    )
-
-    logger.info("Exporting model to %s", output_dir)
-    export_manager.save(output_dir)