Architecture Overview¶

jax2onnx converts JAX functions and Flax modules to ONNX by tracing a JAX program, lowering the resulting JAXPR into onnx_ir, and only converting to ONNX protobuf at the public export boundary. The converter is built around an IR-only core plus plugins that own primitive-specific lowering semantics.

The key architectural boundary is:

Core converter: tracing, value ownership, graph assembly, function scopes, structural optimization, export post-processing, and guardrails.
Plugins: primitive registration, optional tracing-time bindings, abstract evaluation, ONNX operator choices, attributes, layout rules, and testcase metadata.

The core is not completely blind to ONNX operator names: optimizer and post-processing passes deliberately inspect IR operator types for conservative graph cleanup and export compatibility. What must stay out of the core is primitive-specific lowering policy such as how lax.conv_general_dilated, nnx.Conv, or attention layers map their semantics to ONNX nodes.

Design Principles¶

IR first. converter/ and plugins/ operate on onnx_ir, not ONNX protobuf classes. Protobuf serialization happens in the public API facade.
Plugin-owned semantics. A plugin owns the meaning of the JAX primitive or high-level function it registers. The converter dispatches by primitive name.
Explicit value ownership. Every JAX variable must map to a connected ir.Value. Lowering failures are caught at the equation where a value is missing or disconnected.
Scoped graph bodies. ONNX Functions and control-flow bodies are built in child contexts so inputs, constants, symbolic dimensions, and signatures stay explicit.
Conservative graph cleanup. Optimizer passes are IR-only and structural. They may remove redundant graph artifacts, but they must preserve graph outputs and function signatures.
Source-owned precision. JAX model configuration, typed inputs, constants, and explicit conversions determine computation types. Export does not impose a global arithmetic dtype.

Plugin System Guide - plugin authoring workflow and metadata.
ONNX IR Builder Guide - builder conventions, _outputs, initializers, and graph ownership rules.
Expect Graph Reference - structural testcase assertions.
Control-Flow Body Inputs - body input wiring for If/Loop and function-style subgraphs.
IR Optimizer Guide - pass order and optimizer invariants, including conservative ConvTranspose bias-input folding.
ONNX Functions - user-facing behavior of @onnx_function.
Supported Components and ONNX Operator Coverage - generated coverage views derived from plugin/example metadata.

Module Map¶

Module	Responsibility
`jax2onnx/user_interface.py`	Public `to_onnx`, `onnx_function`, return-mode handling, ONNX protobuf/file serialization.
`jax2onnx/converter/conversion_api.py`	Core conversion pipeline: trace, bind inputs/constants, lower JAXPR, build/finalize `onnx_ir.Model`.
`jax2onnx/converter/lowering_dispatch.py`	Equation-by-equation plugin dispatch, primitive call recording, stacktrace/plugin metadata staging.
`jax2onnx/converter/lowering_registry.py`	Original plugin identities and prepared canonical calls, including legacy registration adapters.
`jax2onnx/converter/output_binding.py`	Native input/output assertions and legacy returned-value binding.
`jax2onnx/converter/ir_context.py`	Plugin-facing context: var-to-value map, constants, graph IO, symbolic dim origins, function registry.
`jax2onnx/converter/ir_builder.py`	Thin wrapper around `onnx_ir` tape builder with project-specific naming, initializer, and metadata behavior.
`jax2onnx/converter/function_scope.py`	Child context used to stage ONNX Function bodies.
`jax2onnx/converter/ir_optimizations.py`	IR-only optimizer pass registry.
`jax2onnx/converter/ir_postprocess.py`	Export preparation: loosen intermediate shapes while preserving input/output types and shapes.
`jax2onnx/plugins/plugin_system.py`	Plugin registry, tracing-time patch activation, `PrimitiveLeafPlugin`, and `@onnx_function`.
`jax2onnx/plugins/`	Primitive and framework-specific lowering implementations plus generated-test metadata.

Conversion Pipeline (Detailed)¶

Public to_onnx(...)
    |
    | user_interface.py
    | - normalize input specs, names, input_params, return mode
    v
conversion_api.to_onnx(...)
    |
    | - normalize inputs to ShapeDtypeStruct
    | - create JAX symbolic dims for string dimensions
    | - activate plugin bindings around tracing
    | - trace with jax.make_jaxpr
    v
ClosedJaxpr
    |
    | - bind constvars
    | - bind graph inputs, including optional NCHW bridges
    | - lower equations through prepared registry calls
    | - assert every output var is bound to a connected ir.Value
    | - bind graph outputs, including optional NCHW bridges
    v
Raw onnx_ir graph
    |
    | - attach ONNX Function bodies
    | - run IR optimizer
    | - apply late attribute overrides and compatibility fixes
    | - normalize value shapes and symbolic dims
    v
Precise-shape onnx_ir.Model
    |
    | user_interface.py
    | - loosen intermediate shapes for runtime flexibility
    | - materialize referenced input_params as graph inputs
    | - apply custom input/output names
    | - return ir, proto, or file
    v
Export result

Core Conversion Stages¶

1. Input Normalization and Tracing¶

The public API accepts shape tuples, arrays, jax.ShapeDtypeStruct, and related input specs. conversion_api normalizes them to ShapeDtypeStruct instances. String dimensions such as "B" are converted to JAX symbolic dimensions inside a shared jax.export.SymbolicScope.

Shape-only specs default to FP32, including under jax.enable_x64(). Typed specs retain their declared dtype; if JAX would narrow one in the effective configuration, export raises an actionable error before tracing the model. FP64 preparation therefore includes both typed inputs and a JAX configuration that supports them. The public to_onnx and allclose APIs share this contract; the removed enable_double_precision keyword is diagnosed whenever supplied.

Plugin bindings are active only around tracing. This is where high-level calls such as NNX modules, decorated @onnx_function targets, or plugin-specific JAX wrappers can emit custom primitive names into the ClosedJaxpr. After tracing, patches are restored.

input_params are passed to the traced callable as fixed keyword arguments. If a later ONNX Function call needs one of those values as a dynamic function input, the function plugin records that requirement and the public facade can materialize the referenced parameter as an explicit graph input.

2. Context Setup¶

IRContext owns the active IRBuilder, the JAX-var-to-IR-value map, symbolic dimension origins, function registry, and attribute overrides. Before lowering, the converter:

binds JAXPR constants as initializers in top-level graphs;
emits Constant nodes instead when inside ONNX Function bodies;
creates graph inputs for JAXPR input variables;
inserts NCHW/NHWC bridge transposes only for explicitly requested external IO.

Constants retain their traced dtype and payload, including when shared by consumers or captured in child scopes. Named functions preserve the effective JAX precision context for delayed tracing and function reuse.

3. Equation Lowering¶

lowering_dispatch.lower_jaxpr_with_plugins(...) walks jaxpr.eqns in order. The registry prepares the canonical calls before the equation loop. For each equation it:

reads eqn.primitive.name;
looks up the prepared lowering for that primitive;
verifies all input vars are already bound;
invokes lower(ctx, eqn) -> None;
verifies every non-drop outvar is bound to a value owned by the current graph.

This makes the plugin contract small but strict: a lowering must produce the IR values for the equation outputs and connect them to the graph. Inputs use require_value_for_var, outputs use allocate_value_for_var, and lookup_value_for_var reads existing bindings without allocating. Legacy parameter signatures, returned values, and converter handlers are normalized at registration; they do not add call-form guessing to the equation loop. See the executable plugin pattern.

4. Model Finalization¶

After graph outputs are bound, IRBuilder.to_ir_model(...) creates the onnx_ir.Model. The converter then:

attaches staged ONNX Functions to the model function store;
runs optimize_graph;
applies late attribute overrides and fills missing Concat(axis=0) attributes;
normalizes value shapes so symbolic dimensions are represented as ir.SymbolicDim objects where needed.

Optimizer failures are non-fatal by default and logged. Set JAX2ONNX_STRICT_OPTIMIZER_FAILURES=1 when the optimizer must fail the export.

5. Public Export Post-Processing¶

The public facade calls postprocess_ir_model(...) before returning or serializing. This pass loosens intermediate value shapes while preserving graph input/output shapes, recursively handles control-flow subgraphs and ONNX Function bodies, and preserves their dtypes and constant payloads.

Only after post-processing does the public facade apply custom IO names, materialize referenced input_params, convert to onnx.ModelProto, or write a file.

Plugin Contract¶

Most primitive plugins subclass PrimitiveLeafPlugin and are registered with @register_native_primitive(...). A plugin normally provides:

Metadata: JAX primitive name, user-facing component name, ONNX operator references, generated testcase metadata, and structural checks.
Binding specs: optional tracing-time patches for high-level APIs. If the primitive is already emitted by JAX directly, no patch is needed.
Abstract eval: shape/dtype inference used by JAX tracing.
Lowering: ONNX IR emission through ctx.builder and binding of equation outvars.

Lowerings should follow these rules:

Require inputs with ctx.require_value_for_var(...); missing non-literal bindings are errors. Literals materialize in their current scope.
Allocate or reuse output placeholders with ctx.allocate_value_for_var(...).
Emit nodes through ctx.builder unless a specific onnx_ir feature requires manual ir.Node construction.
Prefer typed op(..., output=...) / op_multi_out(..., outputs=...); keep _outputs sequence-valued in existing dynamic shorthand.
Stamp dtype/shape metadata on produced values when the builder cannot infer it.
Bind outputs with ctx.bind_value_for_var(...). Returned-value conventions remain at the legacy registration boundary.
Never import ONNX protobuf classes from converter or plugin lowering code.

When an ONNX operator needs a different work dtype, the plugin emits explicit consumer-local conversions and restores the source result type at the required boundary. It must preserve canonical input bindings and shared initializers. These adaptations belong to the operation, rather than an exporter-wide precision setting. See the builder guide.

Native registration stores lower(ctx, eqn) -> None directly; only legacy extensions need call-convention discovery or result adaptation. New native plugins use NativeLoweringContext, whose builder exposes the closed IRBuilderAPI protocol. Internal dispatch and scope management use ConverterContext. The legacy LoweringContextProtocol alias and dynamic IRBuilderProtocol retain existing plugin conventions.

ONNX Function Boundaries¶

@onnx_function is user-facing, but its lowering path is part of the converter architecture. Decorating a function or class registers a FunctionPlugin. During tracing, calls bind an onnx_fn::<qualified-name> primitive. During lowering, the plugin:

builds a key from the qualified target name, input aval signatures, and captured/static parameters;
reuses a cached function definition when the key matches;
otherwise opens a FunctionScope child context;
maps parent values to function inputs;
retraces and lowers the callable body inside the child context;
emits constants as Constant nodes because Function bodies cannot own initializers;
attaches the resulting onnx_ir.Function to the parent model and emits a call-site node.

Identical calls share a function body when the function key matches. Shape, dtype, instance, namespace, or capture changes can produce a separate function definition.

FunctionScope.end(outputs) finalizes the child graph once and returns an ir.Function. The registry caches that same object, which model assembly also receives. Inputs, outputs, and nodes belong to its native graph rather than a second project-specific function representation. Finalization clones the body graph to keep its values independent of the construction context and preserves opsets, names, and late attribute overrides.

Control Flow¶

Control-flow primitives are plugins too:

lax.cond lowers to ONNX If.
lax.while_loop lowers to ONNX Loop.
lax.scan currently lowers through ONNX Loop, not the ONNX Scan operator.

Loop bodies use explicit inputs for carried state, sequence values, and captured tensors. If branches have no formal inputs and capture outer values by name. Constants inside those bodies are materialized as Constant nodes. Body shape handling is deliberately conservative, and export post-processing loosens nested body shapes where runtime shape inference needs more flexibility.

See Control-Flow Body Inputs for the detailed body-input rules.

Shapes and Symbolic Dimensions¶

Symbolic dimensions enter through user input specs, are preserved in JAX abstract values, and are tracked in IRContext as origins from a source tensor and axis. Plugins should not cast symbolic dimensions to Python integers.

When an ONNX operator needs a runtime shape value, plugins build it from IR operations such as:

Shape(x)
Gather(shape, axis)
Squeeze or Unsqueeze
Concat
Reshape(x, runtime_shape)

The dim_as_value path uses the recorded symbolic-dim origin to materialize a runtime dimension from the correct tensor instead of guessing from labels alone.

Optimizer Boundary¶

The optimizer runs after raw IR lowering and before late attribute overrides. Its pass order is documented in the IR Optimizer Guide. The graph rewrite reference demonstrates typed outputs, metadata preservation, and native ownership.

Optimizer passes may inspect ONNX op types, but they must remain structural and IR-only. They may fold redundant casts, transposes, reshapes, constants, or dead nodes. They must not encode the semantics of a JAX primitive that belongs in a plugin, and they must preserve graph outputs and ONNX Function signatures.

Structural Tests and Generated Coverage¶

Plugin/example metadata drives both generated tests and user-facing coverage pages. Structural graph checks should live beside the metadata that owns the behavior:

from jax2onnx.plugins._post_check_onnx_graph import expect_graph as EG

Use expect_graph(...) to assert durable structure: operator paths, counts, shapes, attributes, absent operators, unused inputs, and imported ONNX Function bodies when relevant. It is not a general recursive control-flow body walker; Loop/If body contracts should be tested directly when body arity or wiring is the behavior under test.

When lowering structure changes, regenerate a candidate with:

poetry run python scripts/emit_expect_graph.py <testcase>

Then simplify the snippet before storing it in metadata. Generated coverage tables in the user guide should remain generated artifacts, not handwritten architecture claims.

Architectural Guardrails¶

Keep converter and plugin lowering code onnx_ir-only.
Keep primitive semantics in plugins, not in conversion_api.
Use ctx.builder for normal node emission.
Treat graph outputs, function inputs, and symbolic-dim origins as public contracts during rewrites.
Do not rely on private onnx_ir mirrors unless a documented compatibility shim owns the boundary.
Keep plugin activation scoped to tracing and nested function-body tracing.
Preserve deterministic module construction in testcase metadata; use construct_and_call(...) with with_requested_dtype(), with_rng_seed(...), and related placeholder arguments rather than import-time seeding.

Failure Modes¶

Failure	Where it is caught	Typical fix
Plugin binding not active during tracing	Missing primitive-specific JAXPR equation or "no plugin registered" error	Ensure the plugin declares the needed binding spec and tracing runs inside plugin activation.
Missing plugin for primitive	`lowering_dispatch.get_registered_lowering_plugin`	Register a plugin or rewrite the source expression to an already supported primitive.
Plugin reads an unbound input	`assert_eqn_inputs_bound`	Bind constants/inputs earlier or fix nested body input mapping.
Plugin does not bind an output	`assert_eqn_outputs_bound`	Bind each non-drop outvar explicitly. Legacy returned values are handled by the registration adapter.
Plugin binds a disconnected value	`assert_eqn_outputs_bound`	Reuse the pre-allocated output value or emit a node whose output is the bound value.
Function body captures hidden state	Function key / FunctionScope behavior	Make captures explicit as static signature items or dynamic function inputs.
Optimizer removes too much	Focused optimizer regression tests	Make the pass more conservative and preserve outputs/signatures explicitly.

Maintainer Blueprint¶

Add or update a plugin when behavior depends on JAX primitive semantics.
Keep lowering output values connected and metadata stamped.
Add or update generated testcase metadata and focused expect_graph checks.
Use optimizer passes only for shared IR cleanup patterns.
Run focused tests first, then strict docs and broader checks for larger architecture changes.