API Reference

jax2onnx exposes four public entry points:

• to_onnx(...) for export
• @onnx_function for reusable subgraphs
• allclose(...) for JAX-vs-ONNX validation
• allclose_onnxruntime_web(...) for ONNX Runtime Web/WASM parity checks

Common Export Flow

from jax2onnx import to_onnx


model_path = to_onnx(
    fn,
    inputs=[("B", 128)],
    return_mode="file",
    output_path="model.onnx",
)

Use string dimensions such as "B" when you want symbolic dynamic axes.

For direct file export, set return_mode="file" and provide output_path.

For a full export validation checklist, see Validation & Deployment Readiness.

For browser deployment with onnxruntime-web, use the Web export profile:

model_path = to_onnx(
    fn,
    inputs=[("B", 128)],
    return_mode="file",
    output_path="model.web.onnx",
    export_mode="web",
)

export_mode="web" keeps the ONNX graph semantics unchanged, but serializes a single self-contained .onnx file instead of spilling large initializers into a .onnx.data sidecar. That is the easiest artifact shape to serve to onnxruntime-web/wasm.

Parameters To Reach For First

• inputs: Positional input specs, either concrete arrays, ShapeDtypeStruct values, or shape tuples like ("B", 128).
• input_params: Keyword arguments supplied during tracing, normally fixed at export. Values required dynamically by ONNX Functions may become graph inputs; inspect the exported interface before treating a flag as a runtime input.
• return_mode: "proto" for an onnx.ModelProto, "ir" for the intermediate onnx_ir.Model, or "file" to serialize directly to disk.
• export_mode: "standard" for normal serialization, or "web" for single-file browser/WASM artifacts.
• normalization_mode: Export representation for Flax NNX and Linen GroupNorm, and Equinox, Flax NNX, and Linen RMSNorm and LayerNorm. "auto" (default) follows the plugin's current accuracy selection; "prefer_native" requests an eligible standard ONNX operator; and "force_decomposed" emits the explicit primitive graph. See Normalization Representation and Known Limitations for the selection rules and tested accuracy limits.
• inputs_as_nchw / outputs_as_nchw: Adapt the external ONNX interface to NCHW while keeping the traced JAX computation in its original layout.
• input_names / output_names: Apply stable user-facing names after conversion.

Dtypes and Precision

to_onnx(...) exports the computation prepared in JAX. Input types, parameters, explicit casts, and framework computation settings determine its precision. Shape-only inputs such as ("B", 128) default to FP32, even when JAX x64 is enabled. Use arrays or jax.ShapeDtypeStruct to declare other input types.

For FP64, configure the model and typed inputs in an effective JAX x64 context, and use the same context for allclose(...):

import jax
import jax.numpy as jnp
import numpy as np

from jax2onnx import allclose, to_onnx


def configured_model(x):
    return jnp.sin(x)


with jax.enable_x64():
    model_path = to_onnx(
        configured_model,
        inputs=[jax.ShapeDtypeStruct((2, 128), jnp.float64)],
        return_mode="file",
        output_path="model.f64.onnx",
    )
    matches, message = allclose(
        configured_model,
        model_path,
        inputs=[np.ones((2, 128), dtype=np.float64)],
    )
    assert matches, message

The x64 context permits FP64; it does not rewrite explicit FP32 operations or fixed framework dtype settings. Casting existing FP32 weights to FP64 preserves their already-rounded values. A declared input type that JAX tracing would narrow is rejected with a configuration diagnostic.

enable_double_precision has been removed from both to_onnx and allclose. Passing it with any value, including False or None, raises TypeError before export or validation. Remove the argument and prepare the model, inputs, and JAX configuration explicitly.

JAX inputs retain their weak_type promotion behaviour, including when passed through @onnx_function or captured by it. For example, a weak FP32 scalar can combine with an explicitly typed FP16 value in FP16; a strong FP32 operand instead requests FP32 promotion. NumPy inputs remain strongly typed.

lax.dot_general with INT8 or UINT8 operands and an INT32 result uses integer accumulation through MatMulInteger. This also applies to ordinary JAX computations emitted by optional quantization libraries. Preparing or training a quantized model remains a separate step before to_onnx. Conversion from FP16, BF16, FP32 or FP64 to INT8 or UINT8 preserves saturation at the integer limits and maps NaN to zero; integer-to-integer conversions retain wrapping behaviour.

Migration from v0.17.1

In v0.17.1, omitting enable_double_precision or passing False disabled JAX x64 during export; True enabled exporter-wide promotion. Since 0.18.0, export preserves the caller's effective JAX configuration and the types chosen by the prepared model. Removing the old argument alone may therefore change results. To retain an explicitly FP32 computation, configure its parameters and inputs as FP32, and use a caller-owned jax.enable_x64(False) context where needed. For FP64, prepare the model and typed inputs in an enabled x64 context as above; enabling x64 alone does not widen fixed FP32 weights or arithmetic.

allclose also preserves the caller's matrix-precision configuration instead of temporarily forcing jax.default_matmul_precision("float32"). Use matching JAX settings for model preparation, export and validation. NumPy reference inputs whose dtype JAX would narrow are rejected rather than silently changed. Validate representative inputs against your intended JAX computation when migrating; portable ONNX floating-point execution is not a bitwise guarantee.

Normalization Representation

For example, this opts a Fast-Variance GroupNorm into GroupNormalization with opset 21+, while retaining an explicit compatibility artifact for runtimes or inputs where the exported graph should retain the framework reduction layout:

prefer_native_model = to_onnx(
    fn,
    inputs,
    opset=21,
    normalization_mode="prefer_native",
)
force_decomposed_model = to_onnx(
    fn,
    inputs,
    opset=21,
    normalization_mode="force_decomposed",
)

The "auto" policy has two steps: first, choose the implementation with the best reproducible accuracy in a defined test environment and fix its acceptance bounds; then prefer a native operator if it can replace that choice while meeting the same bounds. This is a release qualification policy; to_onnx does not benchmark each model. The current selection is an explicit graph for Flax NNX and Linen GroupNorm, and Equinox, Flax NNX, and Linen RMSNorm and LayerNorm. Fixed comparative bounds currently cover selected float32 LayerNorm cases on ONNX Runtime CPU; equivalent bounds for GroupNorm and RMSNorm have not yet been established.

"prefer_native" requests LayerNormalization at opset 17+, fast-variance GroupNormalization at opset 21+, or RMSNormalization at opset 23+ when the plugin can map the operation and the ONNX builder supports it. Ineligible cases use the explicit graph, including slow-variance GroupNorm, statically empty GroupNorm shapes, and requests below the relevant opset. Linen LayerNorm with use_fast_variance=False traces the original JAX computation in every mode. Native eligibility does not include a per-model accuracy check; The selected native LayerNorm cases are tested against separate, looser limits than "auto", as documented in Known Limitations. Native operators can yield smaller graphs but may have different runtime precision.

"force_decomposed" explicitly chooses the primitive graph at export. Runtime graph optimization happens after export and may rewrite either graph; the tested ONNX Runtime configuration preserves the explicit LayerNorm and RMSNorm paths covered by the tests. See Known Limitations for the fixed LayerNorm acceptance limits, test conditions, and deployment guidance.

Browser/WASM Validation

The generated test harness can optionally validate exported models with onnxruntime-web/wasm in Node.js or Chrome/Chromium:

npm install
JAX2ONNX_VALIDATE_ONNXRUNTIME_WEB=1 poetry run pytest -q tests/primitives/test_nn.py

For the browser runner, add:

npx playwright install chromium
JAX2ONNX_VALIDATE_ONNXRUNTIME_WEB=1 \
JAX2ONNX_ONNXRUNTIME_WEB_RUNNER=chrome \
poetry run pytest -q tests/primitives/test_nn.py

When this flag is enabled, generated tests export with export_mode="web", keep the existing JAX-vs-Python-ONNX-Runtime CPU check, then compare the same ONNX model and inputs against onnxruntime-web/wasm.

For a smaller local smoke run that covers the Quickstart Web model plus representative generated LAX/JAX NumPy examples, run the explicit smoke scripts:

scripts/run_onnxruntime_web_smoke.sh
scripts/run_onnxruntime_web_chrome_smoke.sh

When Web runtime validation is requested through the central repository check runner, it runs the full pytest suite with export_mode="web" and the selected runtime runner:

JAX2ONNX_RUN_ONNXRUNTIME_WEB=1 ./scripts/run_all_checks.sh
JAX2ONNX_RUN_ONNXRUNTIME_WEB_CHROME=1 ./scripts/run_all_checks.sh

For browser loading code, validation helpers, CI usage, and troubleshooting, see Browser/WASM Deployment.

Use @onnx_function when repeated callables should become reusable ONNX functions, and allclose(...) when you want a quick numerical check against ONNX Runtime after export.

jax2onnx.user_interface

User-facing API for converting JAX functions and models to ONNX.

This module provides the primary interface for exporting JAX/Flax models to the ONNX format. It supports dynamic shapes, runtime parameters, file-oriented export profiles, and numerical validation against ONNX Runtime.

Key Functions:

• to_onnx: Convert a JAX function or Flax module to an ONNX model.
• onnx_function: Decorator to mark a function or class as an ONNX function node.
• allclose: Validate numerical equivalence between JAX and ONNX Runtime outputs.
• allclose_onnxruntime_web: Validate ONNX Runtime Web/WASM output against Python ONNX Runtime CPU output.

Example

import jax.numpy as jnp

from jax2onnx import to_onnx


def my_model(x):
    return jnp.sin(x)


to_onnx(
    my_model,
    inputs=[("B", 10)],
    return_mode="file",
    output_path="model.onnx",
)

to_onnx(
    my_model,
    inputs=[("B", 10)],
    return_mode="file",
    output_path="model.web.onnx",
    export_mode="web",
)

allclose(fn, onnx_model_path, inputs, input_params=None, rtol=0.001, atol=1e-05, *, inputs_as_nchw=None, outputs_as_nchw=None, **removed_options)

Checks if JAX and ONNX Runtime outputs remain numerically close.

Parameters:

Name	Type	Description	Default
fn	Callable[..., Any]	JAX callable to compare against the exported ONNX model.	required
onnx_model_path	str	Path to a serialized model that ORT can execute.	required
inputs	Sequence[Any]	Concrete input arrays (or shape tuples, which will be sampled).	required
input_params	Optional[Dict[str, Any]]	Optional keyword arguments applied to both call sites.	None
rtol	float	Relative tolerance for floating-point comparisons.	0.001
atol	float	Absolute tolerance for floating-point comparisons.	1e-05

Returns:

Type	Description
bool	(is_match, message) where is_match indicates success and message
str	provides context when a mismatch occurs.

Example

import jax.numpy as jnp

from jax2onnx import allclose, to_onnx


def my_func(x):
    return jnp.sin(x)


model_path = to_onnx(
    my_func,
    inputs=[("B", 10)],
    return_mode="file",
    output_path="my_model.onnx",
)

validation_inputs = [(5, 10)]
is_match, msg = allclose(
    my_func,
    model_path,
    inputs=validation_inputs,
    atol=1e-5,
)

assert is_match, f"Validation failed: {msg}"

Source code in jax2onnx/user_interface.py

def allclose(
    fn: Callable[..., Any],
    onnx_model_path: str,
    inputs: Sequence[Any],
    input_params: Optional[Dict[str, Any]] = None,
    rtol: float = 1e-3,
    atol: float = 1e-5,
    *,
    inputs_as_nchw: Optional[Sequence[int]] = None,
    outputs_as_nchw: Optional[Sequence[int]] = None,
    **removed_options: Never,
) -> Tuple[bool, str]:
    """
    Checks if JAX and ONNX Runtime outputs remain numerically close.

    Args:
        fn: JAX callable to compare against the exported ONNX model.
        onnx_model_path: Path to a serialized model that ORT can execute.
        inputs: Concrete input arrays (or shape tuples, which will be sampled).
        input_params: Optional keyword arguments applied to both call sites.
        rtol: Relative tolerance for floating-point comparisons.
        atol: Absolute tolerance for floating-point comparisons.

    Returns:
        `(is_match, message)` where `is_match` indicates success and `message`
        provides context when a mismatch occurs.

    Example:
        ```python
        import jax.numpy as jnp

        from jax2onnx import allclose, to_onnx


        def my_func(x):
            return jnp.sin(x)


        model_path = to_onnx(
            my_func,
            inputs=[("B", 10)],
            return_mode="file",
            output_path="my_model.onnx",
        )

        validation_inputs = [(5, 10)]
        is_match, msg = allclose(
            my_func,
            model_path,
            inputs=validation_inputs,
            atol=1e-5,
        )

        assert is_match, f"Validation failed: {msg}"
        ```
    """

    if "enable_double_precision" in removed_options:
        raise TypeError(
            "enable_double_precision was removed; use the same JAX context as export"
        )
    if removed_options:
        unexpected = next(iter(removed_options))
        raise TypeError(f"allclose() got an unexpected keyword argument {unexpected!r}")

    _LOGGER.info(
        "Comparing JAX and ONNX outputs (path=%s, rtol=%s, atol=%s)",
        onnx_model_path,
        rtol,
        atol,
    )

    xs = _validation_inputs_to_arrays(inputs)

    params = dict(input_params or {})
    with preserve_precision_globals():
        return _run_allclose(
            fn,
            onnx_model_path,
            xs,
            params,
            rtol=rtol,
            atol=atol,
            inputs_as_nchw=inputs_as_nchw,
            outputs_as_nchw=outputs_as_nchw,
        )

allclose_onnxruntime_web(onnx_model_path, inputs, input_params=None, rtol=0.001, atol=1e-05, *, inputs_as_nchw=None, node_command='node', runner=None)

Validate that onnxruntime-web/wasm matches Python ONNX Runtime CPU output.

This helper is intended for deployment-readiness checks. It executes the serialized model once with Python onnxruntime on CPUExecutionProvider, then executes the same model and feeds through onnxruntime-web/wasm. The default runner is Node.js; pass runner="chrome" or set JAX2ONNX_ONNXRUNTIME_WEB_RUNNER=chrome to validate in a real browser via Playwright/Chromium.

Parameters:

Name	Type	Description	Default
onnx_model_path	str	Path to a serialized ONNX model, typically produced with to_onnx(..., return_mode="file", export_mode="web").	required
inputs	Sequence[Any]	Concrete input arrays (or shape tuples, which will be sampled).	required
input_params	Optional[Dict[str, Any]]	Optional named ONNX inputs that correspond to runtime parameters materialized during export.	None
rtol	float	Relative tolerance for output comparison.	0.001
atol	float	Absolute tolerance for output comparison.	1e-05
inputs_as_nchw	Optional[Sequence[int]]	Optional sequence of input indices that should be transposed from NHWC validation arrays to NCHW ONNX feeds.	None
node_command	str	Node.js executable used to run the Web/WASM validator.	'node'
runner	Optional[str]	Optional runtime selector. Use "node" for Node.js WASM validation or "chrome"/"browser" for Playwright Chromium.	None

Returns:

Type	Description
bool	(is_match, message) where is_match indicates success and message
str	contains the Web runner result or failure details.

Example

import numpy as np

from jax2onnx import allclose_onnxruntime_web


is_match, message = allclose_onnxruntime_web(
    "model.web.onnx",
    inputs=[np.zeros((2, 10), dtype=np.float32)],
    rtol=1e-5,
    atol=1e-5,
    runner="node",
)

assert is_match, message

Source code in jax2onnx/user_interface.py

def allclose_onnxruntime_web(
    onnx_model_path: str,
    inputs: Sequence[Any],
    input_params: Optional[Dict[str, Any]] = None,
    rtol: float = 1e-3,
    atol: float = 1e-5,
    *,
    inputs_as_nchw: Optional[Sequence[int]] = None,
    node_command: str = "node",
    runner: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Validate that `onnxruntime-web/wasm` matches Python ONNX Runtime CPU output.

    This helper is intended for deployment-readiness checks. It executes the
    serialized model once with Python `onnxruntime` on `CPUExecutionProvider`,
    then executes the same model and feeds through `onnxruntime-web/wasm`.
    The default runner is Node.js; pass `runner="chrome"` or set
    `JAX2ONNX_ONNXRUNTIME_WEB_RUNNER=chrome` to validate in a real browser via
    Playwright/Chromium.

    Args:
        onnx_model_path: Path to a serialized ONNX model, typically produced with
            `to_onnx(..., return_mode="file", export_mode="web")`.
        inputs: Concrete input arrays (or shape tuples, which will be sampled).
        input_params: Optional named ONNX inputs that correspond to runtime
            parameters materialized during export.
        rtol: Relative tolerance for output comparison.
        atol: Absolute tolerance for output comparison.
        inputs_as_nchw: Optional sequence of input indices that should be
            transposed from NHWC validation arrays to NCHW ONNX feeds.
        node_command: Node.js executable used to run the Web/WASM validator.
        runner: Optional runtime selector. Use `"node"` for Node.js WASM
            validation or `"chrome"`/`"browser"` for Playwright Chromium.

    Returns:
        `(is_match, message)` where `is_match` indicates success and `message`
        contains the Web runner result or failure details.

    Example:
        ```python
        import numpy as np

        from jax2onnx import allclose_onnxruntime_web


        is_match, message = allclose_onnxruntime_web(
            "model.web.onnx",
            inputs=[np.zeros((2, 10), dtype=np.float32)],
            rtol=1e-5,
            atol=1e-5,
            runner="node",
        )

        assert is_match, message
        ```
    """

    xs = _validation_inputs_to_arrays(inputs)
    params = dict(input_params or {})

    try:
        ort = cast(Any, importlib.import_module("onnxruntime"))
    except ImportError as exc:
        return False, f"onnxruntime is required for CPU reference execution: {exc}"

    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    sess_options.enable_mem_pattern = False
    sess_options.intra_op_num_threads = 1
    sess_options.inter_op_num_threads = 1

    try:
        session = ort.InferenceSession(
            onnx_model_path,
            sess_options=sess_options,
            providers=["CPUExecutionProvider"],
        )
    except Exception as exc:  # pragma: no cover - defensive aid
        return False, f"Failed to create CPU ONNX Runtime session: {exc}"

    ort_xs = list(xs)
    if inputs_as_nchw:
        nhwc_to_nchw = [0, 3, 1, 2]
        for idx in inputs_as_nchw:
            if 0 <= idx < len(ort_xs):
                val = ort_xs[idx]
                if hasattr(val, "ndim") and val.ndim == 4:
                    ort_xs[idx] = np.transpose(val, nhwc_to_nchw)

    try:
        ort_inputs = _build_ort_inputs(session, ort_xs, params)
        expected_outputs = session.run(None, ort_inputs)
    except Exception as exc:  # pragma: no cover - defensive aid
        return False, f"Failed to run CPU ONNX Runtime reference: {exc}"

    output_names = [output_meta.name for output_meta in session.get_outputs()]
    if len(output_names) != len(expected_outputs):
        return (
            False,
            f"Output metadata mismatch (names={len(output_names)} outputs={len(expected_outputs)})",
        )

    try:
        spec = {
            "modelPath": str(Path(onnx_model_path).resolve()),
            "rtol": float(rtol),
            "atol": float(atol),
            "inputs": [
                _array_to_onnxruntime_web_tensor(name, value)
                for name, value in ort_inputs.items()
            ],
            "outputs": [
                _array_to_onnxruntime_web_tensor(name, value)
                for name, value in zip(output_names, expected_outputs, strict=True)
            ],
        }
    except TypeError as exc:
        return False, str(exc)

    configured_runner = (
        runner
        if runner is not None
        else os.getenv("JAX2ONNX_ONNXRUNTIME_WEB_RUNNER", "node")
    )
    runner_name = configured_runner.strip().lower()
    if runner_name == "node":
        script_name = "validate_onnxruntime_web.mjs"
    elif runner_name in {"browser", "chrome"}:
        script_name = "validate_onnxruntime_web_chrome.mjs"
        runner_name = "chrome"
    else:
        return (
            False,
            "Unsupported onnxruntime-web runner "
            f"'{runner_name}'. Use 'node' or 'chrome'.",
        )

    script_path = Path(__file__).resolve().parent.parent / "scripts" / script_name
    if not script_path.is_file():
        return (
            False,
            f"onnxruntime-web/{runner_name} validation script not found: {script_path}",
        )

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".json", encoding="utf-8", delete=False
    ) as tmp:
        spec_path = Path(tmp.name)
        json.dump(spec, tmp, allow_nan=False)

    try:
        completed = subprocess.run(
            [node_command, str(script_path), str(spec_path)],
            cwd=str(script_path.parent.parent),
            text=True,
            capture_output=True,
            check=False,
            env=os.environ.copy(),
        )
    except FileNotFoundError:
        return False, f"Node.js command not found: {node_command}"
    finally:
        try:
            spec_path.unlink()
        except OSError:
            pass

    if completed.returncode != 0:
        stderr = completed.stderr.strip()
        stdout = completed.stdout.strip()
        details = stderr or stdout or f"exit code {completed.returncode}"
        return False, f"onnxruntime-web/{runner_name} validation failed: {details}"

    summary = (
        completed.stdout.strip() or f"onnxruntime-web/{runner_name} outputs match."
    )
    return True, summary

onnx_function(target=None, *, unique=False, namespace=None, name=None, type=None)

onnx_function(
    target: None = ...,
    *,
    unique: bool = ...,
    namespace: Optional[str] = ...,
    name: Optional[str] = ...,
    type: Optional[str] = ...
) -> OnnxFunctionDecorator

onnx_function(
    target: OnnxFunctionTarget,
    *,
    unique: bool = ...,
    namespace: Optional[str] = ...,
    name: Optional[str] = ...,
    type: Optional[str] = ...
) -> OnnxFunctionTarget

Decorator to mark a function or class as an ONNX function.

This decorator is used to indicate that a function or class should be converted to an ONNX function node when included in a model. It allows the function to be traced and exported as a reusable component with its own namespace in the ONNX graph.

Parameters:

Name	Type	Description	Default
target	Optional[OnnxFunctionTarget]	The target function or class to decorate. When omitted, the decorator must be called with parentheses.	None
unique	bool	If True, reuse a single ONNX Function definition for all call sites that share the same callable type and captured parameters.	False
namespace	Optional[str]	Custom domain prefix for the emitted FunctionProto. Defaults to "custom" when omitted.	None
name	Optional[str]	Optional human-readable base name for the ONNX function. When set, this overrides the callable's Python name for the function op_type and FunctionProto name; the domain still derives from namespace.	None
type	Optional[str]	Alias for name; preferred keyword for setting the function op_type/display name in ONNX. If both name and type are supplied, type takes precedence.	None

Raises:

Type	Description
ValueError	If the same target is decorated again with a conflicting namespace or display-name override.

Returns:

Type	Description
Union[OnnxFunctionTarget, OnnxFunctionDecorator]	The decorated function or class with ONNX function capabilities.

Example

import jax.numpy as jnp
from flax import nnx

from jax2onnx import onnx_function


@onnx_function
def my_custom_op(x, y):
    return jnp.sin(x) * y


@onnx_function
class MLPBlock(nnx.Module):
    def __init__(self, features, rngs):
        self.dense = nnx.Linear(features, rngs=rngs)
        self.activation = nnx.relu

    def __call__(self, x):
        return self.activation(self.dense(x))

Source code in jax2onnx/user_interface.py

def onnx_function(
    target: Optional[OnnxFunctionTarget] = None,
    *,
    unique: bool = False,
    namespace: Optional[str] = None,
    name: Optional[str] = None,
    type: Optional[str] = None,  # noqa: A002 - user-facing keyword
) -> Union[OnnxFunctionTarget, OnnxFunctionDecorator]:
    """
    Decorator to mark a function or class as an ONNX function.

    This decorator is used to indicate that a function or class should be converted to
    an ONNX function node when included in a model. It allows the function to be traced
    and exported as a reusable component with its own namespace in the ONNX graph.

    Args:
        target: The target function or class to decorate. When omitted, the decorator
            must be called with parentheses.
        unique: If True, reuse a single ONNX Function definition for all call sites
            that share the same callable type and captured parameters.
        namespace: Custom domain prefix for the emitted FunctionProto. Defaults to
            ``"custom"`` when omitted.
        name: Optional human-readable base name for the ONNX function. When set,
            this overrides the callable's Python name for the function `op_type`
            and FunctionProto name; the domain still derives from ``namespace``.
        type: Alias for ``name``; preferred keyword for setting the function
            `op_type`/display name in ONNX. If both ``name`` and ``type`` are
            supplied, ``type`` takes precedence.

    Raises:
        ValueError: If the same target is decorated again with a conflicting
            namespace or display-name override.

    Returns:
        The decorated function or class with ONNX function capabilities.

    Example:
        ```python
        import jax.numpy as jnp
        from flax import nnx

        from jax2onnx import onnx_function


        @onnx_function
        def my_custom_op(x, y):
            return jnp.sin(x) * y


        @onnx_function
        class MLPBlock(nnx.Module):
            def __init__(self, features, rngs):
                self.dense = nnx.Linear(features, rngs=rngs)
                self.activation = nnx.relu

            def __call__(self, x):
                return self.activation(self.dense(x))
        ```
    """

    return cast(
        Union[OnnxFunctionTarget, OnnxFunctionDecorator],
        onnx_function_impl(
            target, unique=unique, namespace=namespace, name=name, type=type
        ),
    )

to_onnx(fn, inputs, input_params=None, model_name='jax_model', opset=23, *, record_primitive_calls_file=None, return_mode='proto', output_path=None, inputs_as_nchw=None, outputs_as_nchw=None, input_names=None, output_names=None, export_mode='standard', normalization_mode='auto', **removed_options)

to_onnx(
    fn: Callable[..., Any],
    inputs: Sequence[UserInputSpec],
    input_params: Optional[Mapping[str, object]] = ...,
    model_name: str = ...,
    opset: int = ...,
    *,
    record_primitive_calls_file: Optional[str] = ...,
    return_mode: Literal["proto"] = ...,
    output_path: None = ...,
    inputs_as_nchw: Optional[Sequence[int]] = ...,
    outputs_as_nchw: Optional[Sequence[int]] = ...,
    input_names: Optional[Sequence[str]] = ...,
    output_names: Optional[Sequence[str]] = ...,
    export_mode: ExportMode = ...,
    normalization_mode: NormalizationMode = ...
) -> onnx.ModelProto

to_onnx(
    fn: Callable[..., Any],
    inputs: Sequence[UserInputSpec],
    input_params: Optional[Mapping[str, object]] = ...,
    model_name: str = ...,
    opset: int = ...,
    *,
    record_primitive_calls_file: Optional[str] = ...,
    return_mode: Literal["ir"],
    output_path: Optional[PathLikeStr] = ...,
    inputs_as_nchw: Optional[Sequence[int]] = ...,
    outputs_as_nchw: Optional[Sequence[int]] = ...,
    input_names: Optional[Sequence[str]] = ...,
    output_names: Optional[Sequence[str]] = ...,
    export_mode: ExportMode = ...,
    normalization_mode: NormalizationMode = ...
) -> ir.Model

to_onnx(
    fn: Callable[..., Any],
    inputs: Sequence[UserInputSpec],
    input_params: Optional[Mapping[str, object]] = ...,
    model_name: str = ...,
    opset: int = ...,
    *,
    record_primitive_calls_file: Optional[str] = ...,
    return_mode: Literal["file"],
    output_path: PathLikeStr,
    inputs_as_nchw: Optional[Sequence[int]] = ...,
    outputs_as_nchw: Optional[Sequence[int]] = ...,
    input_names: Optional[Sequence[str]] = ...,
    output_names: Optional[Sequence[str]] = ...,
    export_mode: ExportMode = ...,
    normalization_mode: NormalizationMode = ...
) -> str

Converts a JAX function or model into an ONNX model.

This function serves as the main entry point for converting JAX/Flax models to ONNX format. It supports dynamic shapes and additional keyword configuration. Precision follows the prepared JAX computation and caller configuration. Shape-only inputs default to FP32; use typed inputs to request other dtypes.

Parameters:

Name	Type	Description	Default
fn	Callable[..., Any]	The JAX function or Flax module to convert.	required
inputs	Sequence[UserInputSpec]	Sequence of input specifications. Each entry may be: * a jax.ShapeDtypeStruct (or jax.core.ShapedArray); * any array-like object exposing .shape and .dtype (e.g. jax.Array, np.ndarray); * a tuple/list of ints/strs describing the desired shape. Example: [('B', 128), (1, 10)] implies two inputs, the first with a dynamic batch dimension 'B' and fixed size 128.	required
input_params	Optional[Mapping[str, object]]	Keyword arguments supplied when tracing the callable, such as "deterministic" flags. These are normally fixed at export. Referenced values required by ONNX Functions may become graph inputs; inspect the exported interface before treating a flag as runtime input.	None
model_name	str	Name to give the ONNX model. Defaults to "jax_model".	'jax_model'
opset	int	ONNX opset version to target. Defaults to 23.	23
record_primitive_calls_file	Optional[str]	Optional path to a file. If provided, details of each JAX primitive encountered during conversion will be recorded to this file. This log can be used by developers to manually create new test cases. Defaults to None (disabled).	None
return_mode	ReturnMode	Output mode. "proto" (default) returns an ONNX ModelProto, "ir" returns the intermediate onnx_ir.Model, and "file" serialises directly to disk.	'proto'
output_path	Optional[PathLikeStr]	Destination path (str or PathLike) required when return_mode is "file". Ignored otherwise.	None
inputs_as_nchw	Optional[Sequence[int]]	Optional sequence of input indices (0-based) that should be treated as NCHW layout. If specified for an input, jax2onnx assumes the external input is NCHW and will automatically transpose it to NHWC before feeding it to the JAX graph (which typically expects NHWC for images). This allows the exported ONNX model to accept NCHW inputs while preserving correct graph semantics.	None
outputs_as_nchw	Optional[Sequence[int]]	Optional sequence of output indices (0-based) that should be treated as NCHW layout. If specified for an output, jax2onnx assumes the external output should be NCHW and will automatically transpose the NHWC output derived from JAX graph to NCHW before returning it.	None
input_names	Optional[Sequence[str]]	Optional sequence of names for positional model inputs. When provided, names are assigned by positional index (0-based) after conversion. These names do not apply to entries supplied through input_params.	None
output_names	Optional[Sequence[str]]	Optional sequence of names for model outputs (0-based). Names are assigned after conversion in output order.	None
export_mode	ExportMode	Serialization profile. "standard" preserves the existing file behavior, spilling large initializers into .onnx.data sidecars when needed. "web" writes a single self-contained .onnx file for browser/WASM deployment via onnxruntime-web. This only affects return_mode="file"; "proto" and "ir" return values are unchanged.	'standard'
normalization_mode	NormalizationMode	Export policy for normalization plugins that offer both native ONNX operators and explicit graphs. "auto" (default) follows a two-step policy: choose the most reproducibly accurate implementation in a defined test environment and fix its acceptance bounds; then permit a native replacement only if it meets those bounds. Export does not benchmark each model. The current selection is explicit for Flax NNX and Linen GroupNorm, and Equinox, Flax NNX, and Linen RMSNorm and LayerNorm. Fixed comparative bounds currently cover selected LayerNorm cases, not GroupNorm or RMSNorm. "prefer_native" requests a native operator when eligible (LayerNormalization from opset 17, fast-variance GroupNormalization from 21, RMSNormalization from 23), falling back to the explicit graph otherwise. Eligibility is determined by the plugin and opset, without a per-model accuracy check. Linen LayerNorm with use_fast_variance=False traces JAX in every mode. "force_decomposed" selects the explicit graph. Runtime graph optimization is separate from the exported representation. See the user guide for tested limits.	'auto'

Returns:

Type	Description
Union[ModelProto, Model, str]	If return_mode="proto" (default): Returns an onnx.ModelProto object.
Union[ModelProto, Model, str]	If return_mode="ir": Returns an onnx_ir.Model object (intermediate representation).
Union[ModelProto, Model, str]	If return_mode="file": Returns the string path to the saved file.

Raises:

Type	Description
ValueError	If return_mode is "file" but output_path is not provided.
ValueError	If return_mode is invalid.
TypeError	If input_params keys are not strings.

Example

import jax.numpy as jnp

from jax2onnx import to_onnx


def linear(x, w, b):
    return jnp.dot(x, w) + b


input_specs = [
    ("B", 32),  # x: [Batch, 32]
    (32, 10),  # w: [32, 10]
    (10,),  # b: [10]
]

to_onnx(
    linear,
    inputs=input_specs,
    model_name="linear_model",
    return_mode="file",
    output_path="linear_model.onnx",
)

Source code in jax2onnx/user_interface.py

def to_onnx(
    fn: Callable[..., Any],
    inputs: Sequence[UserInputSpec],
    input_params: Optional[Mapping[str, object]] = None,
    model_name: str = "jax_model",
    opset: int = 23,
    *,  # All arguments after this must be keyword-only
    record_primitive_calls_file: Optional[str] = None,
    return_mode: ReturnMode = "proto",
    output_path: Optional[PathLikeStr] = None,
    inputs_as_nchw: Optional[Sequence[int]] = None,
    outputs_as_nchw: Optional[Sequence[int]] = None,
    input_names: Optional[Sequence[str]] = None,
    output_names: Optional[Sequence[str]] = None,
    export_mode: ExportMode = "standard",
    normalization_mode: NormalizationMode = "auto",
    **removed_options: Never,
) -> Union[onnx.ModelProto, ir.Model, str]:
    """
    Converts a JAX function or model into an ONNX model.

    This function serves as the main entry point for converting JAX/Flax models to ONNX format.
    It supports dynamic shapes and additional keyword configuration.
    Precision follows the prepared JAX computation and caller configuration.
    Shape-only inputs default to FP32; use typed inputs to request other dtypes.

    Args:
        fn: The JAX function or Flax module to convert.
        inputs: Sequence of input specifications. Each entry may be:
            * a `jax.ShapeDtypeStruct` (or `jax.core.ShapedArray`);
            * any array-like object exposing `.shape` and `.dtype`
              (e.g. `jax.Array`, `np.ndarray`);
            * a tuple/list of ints/strs describing the desired shape.
              Example: `[('B', 128), (1, 10)]` implies two inputs, the first with a dynamic batch dimension 'B' and fixed size 128.
        input_params: Keyword arguments supplied when tracing the callable,
            such as `"deterministic"` flags. These are normally fixed at export.
            Referenced values required by ONNX Functions may become graph inputs;
            inspect the exported interface before treating a flag as runtime input.
        model_name: Name to give the ONNX model. Defaults to "jax_model".
        opset: ONNX opset version to target. Defaults to 23.
        record_primitive_calls_file: Optional path to a file. If provided,
            details of each JAX primitive encountered during conversion will be
            recorded to this file. This log can be used by developers to manually
            create new test cases. Defaults to None (disabled).
        return_mode: Output mode. `"proto"` (default) returns an ONNX ModelProto,
            `"ir"` returns the intermediate onnx_ir.Model, and `"file"`
            serialises directly to disk.
        output_path: Destination path (str or PathLike) required when `return_mode` is
            `"file"`. Ignored otherwise.
        inputs_as_nchw: Optional sequence of input indices (0-based) that should be treated as NCHW layout.
            If specified for an input, jax2onnx assumes the external input is NCHW and will automatically
            transpose it to NHWC before feeding it to the JAX graph (which typically expects NHWC for images).
            This allows the exported ONNX model to accept NCHW inputs while preserving correct graph semantics.
        outputs_as_nchw: Optional sequence of output indices (0-based) that should be treated as NCHW layout.
            If specified for an output, jax2onnx assumes the external output should be NCHW and will automatically
            transpose the NHWC output derived from JAX graph to NCHW before returning it.
        input_names: Optional sequence of names for positional model inputs.
            When provided, names are assigned by positional index (0-based) after conversion.
            These names do not apply to entries supplied through `input_params`.
        output_names: Optional sequence of names for model outputs (0-based).
            Names are assigned after conversion in output order.
        export_mode: Serialization profile. `"standard"` preserves the existing
            file behavior, spilling large initializers into `.onnx.data` sidecars
            when needed. `"web"` writes a single self-contained `.onnx` file for
            browser/WASM deployment via `onnxruntime-web`. This only affects
            `return_mode="file"`; `"proto"` and `"ir"` return values are unchanged.
        normalization_mode: Export policy for normalization plugins that offer
            both native ONNX operators and explicit graphs. `"auto"` (default)
            follows a two-step policy: choose the most reproducibly accurate
            implementation in a defined test environment and fix its acceptance
            bounds; then permit a native replacement only if it meets those
            bounds. Export does not benchmark each model. The current selection
            is explicit for Flax NNX and Linen GroupNorm, and Equinox, Flax
            NNX, and Linen RMSNorm and LayerNorm. Fixed
            comparative bounds currently cover selected LayerNorm cases,
            not GroupNorm or RMSNorm. `"prefer_native"` requests a native operator
            when eligible (`LayerNormalization` from opset 17, fast-variance
            `GroupNormalization` from 21, `RMSNormalization` from 23), falling
            back to the explicit graph otherwise. Eligibility is determined
            by the plugin and opset, without a per-model accuracy check. Linen
            LayerNorm with use_fast_variance=False traces JAX in every mode.
            `"force_decomposed"` selects the explicit graph. Runtime graph optimization is separate from
            the exported representation. See the user guide for tested limits.

    Returns:
        * If `return_mode="proto"` (default): Returns an `onnx.ModelProto` object.
        * If `return_mode="ir"`: Returns an `onnx_ir.Model` object (intermediate representation).
        * If `return_mode="file"`: Returns the string path to the saved file.

    Raises:
        ValueError: If `return_mode` is "file" but `output_path` is not provided.
        ValueError: If `return_mode` is invalid.
        TypeError: If `input_params` keys are not strings.

    Example:
        ```python
        import jax.numpy as jnp

        from jax2onnx import to_onnx


        def linear(x, w, b):
            return jnp.dot(x, w) + b


        input_specs = [
            ("B", 32),  # x: [Batch, 32]
            (32, 10),  # w: [32, 10]
            (10,),  # b: [10]
        ]

        to_onnx(
            linear,
            inputs=input_specs,
            model_name="linear_model",
            return_mode="file",
            output_path="linear_model.onnx",
        )
        ```
    """

    if "enable_double_precision" in removed_options:
        raise TypeError(
            "enable_double_precision was removed; configure the JAX model and typed inputs instead"
        )
    if removed_options:
        unexpected = next(iter(removed_options))
        raise TypeError(f"to_onnx() got an unexpected keyword argument {unexpected!r}")

    _LOGGER.info(
        f"Converting JAX function to ONNX model with parameters: "
        f"model_name={model_name}, opset={opset}, input_shapes={inputs}, "
        f"input_params={input_params}, "
        f"record_primitive_calls_file={record_primitive_calls_file}, "
        f"return_mode={return_mode}, output_path={output_path}, "
        f"inputs_as_nchw={inputs_as_nchw}, outputs_as_nchw={outputs_as_nchw}, "
        f"input_names={input_names}, output_names={output_names}, "
        f"export_mode={export_mode}, normalization_mode={normalization_mode}"
    )

    # Determine the nature of the 'inputs' argument to prepare for to_onnx_impl
    normalized_mode = _normalize_return_mode(return_mode)
    normalized_export_mode = _normalize_export_mode(export_mode)
    normalized_normalization_mode = _normalize_normalization_mode(normalization_mode)

    file_path: Optional[str] = None
    if normalized_mode == "file":
        if output_path is None:
            raise ValueError(
                "`output_path` must be provided when return_mode is 'file'."
            )
        file_path = os.fspath(output_path)

    normalized_inputs: List[InputSpec] = []
    if inputs:
        normalized_inputs = _normalize_input_specs(inputs)

    normalized_input_names = _normalize_io_names(input_names, kind="input_names")
    normalized_output_names = _normalize_io_names(output_names, kind="output_names")

    if normalized_input_names is not None and len(normalized_input_names) != len(
        normalized_inputs
    ):
        raise ValueError(
            f"input_names length ({len(normalized_input_names)}) must match "
            f"the number of positional inputs ({len(normalized_inputs)})."
        )

    param_map: Dict[str, object] = {}
    if input_params:
        for key, value in input_params.items():
            if not isinstance(key, str):
                raise TypeError(
                    "input_params must use string keys; "
                    f"received key of type {type(key)}."
                )
            param_map[key] = value

    if normalized_input_names is not None:
        collisions = sorted(set(normalized_input_names).intersection(param_map.keys()))
        if collisions:
            names = ", ".join(collisions)
            raise ValueError(
                f"input_names collide with names reserved by input_params: {names}."
            )

    if normalized_output_names is not None:
        collisions = sorted(set(normalized_output_names).intersection(param_map.keys()))
        if collisions:
            names = ", ".join(collisions)
            raise ValueError(
                f"output_names collide with names reserved by input_params: {names}."
            )

    with preserve_precision_globals():
        result = to_onnx_impl(
            fn=fn,
            inputs=normalized_inputs,
            input_params=param_map,
            model_name=model_name,
            opset=opset,
            record_primitive_calls_file=record_primitive_calls_file,
            protective_clone=(normalized_mode == "ir"),
            inputs_as_nchw=inputs_as_nchw,
            outputs_as_nchw=outputs_as_nchw,
            input_names=normalized_input_names,
            output_names=normalized_output_names,
            normalization_mode=normalized_normalization_mode,
        )
        postprocess_ir_model(result)

    def _save_model_proto(
        model_proto: onnx.ModelProto,
        dest: str,
        *,
        mode: ExportMode,
        external_threshold: int = 1_048_576,  # 1 MB default before spilling to .data
    ) -> str:
        dest_dir = os.path.dirname(dest)
        if dest_dir:
            os.makedirs(dest_dir, exist_ok=True)
        data_location = os.path.basename(dest) + ".data"
        data_path = os.path.join(dest_dir or ".", data_location)

        if mode == "web":
            onnx.save_model(model_proto, dest, save_as_external_data=False)
            # A previous standard export to the same path may have left a sidecar.
            try:
                if os.path.exists(data_path):
                    os.remove(data_path)
            except OSError:
                pass
            return dest

        onnx.save_model(
            model_proto,
            dest,
            save_as_external_data=True,
            all_tensors_to_one_file=True,
            location=data_location,
            size_threshold=external_threshold,
            convert_attribute=False,
        )
        # Only keep the .data sidecar if the export actually referenced external data.
        if not any(init.external_data for init in model_proto.graph.initializer):
            # No external payloads; remove an empty sidecar if one was produced.
            try:
                if os.path.exists(data_path) and os.path.getsize(data_path) == 0:
                    os.remove(data_path)
            except OSError:
                pass
        return dest

    _materialize_input_params_on_ir(result, param_map)
    _apply_custom_io_names_on_ir(
        result,
        input_names=normalized_input_names,
        output_names=normalized_output_names,
        positional_input_count=len(normalized_inputs),
    )
    if normalized_mode == "ir":
        return result

    model_proto = ir.to_proto(result)
    if normalized_mode == "file":
        assert file_path is not None
        return _save_model_proto(model_proto, file_path, mode=normalized_export_mode)
    return model_proto