jax2onnx 0.4.2

export JAX to ONNX - focus on flax nnx

jax2onnx 🌟

jax2onnx converts your JAX/Flax functions directly into the ONNX format.

✨ Key Features

• Simple API
  Convert any JAX/Flax model to ONNX using to_onnx(...) or save_onnx(...)

• Model structure preserved
  With @onnx_function, submodules appear as named functions in the ONNX graph (e.g. in Netron). Useful for readability and reuse.

• Dynamic input support
  Use abstract dimensions like 'B' or pass scalars as runtime inputs. Models stay flexible without retracing.

• Plugin-based extensibility
  Add support for new primitives by writing small, local plugins.

• Netron-friendly outputs
  All generated ONNX graphs include shape/type annotations and are structured for clear visualization.

---

🚀 Quickstart

Convert your JAX callable to ONNX in just a few lines:

import onnx
from flax import nnx
from jax2onnx import to_onnx

# Define a simple MLP (from Flax docs)
class MLP(nnx.Module):
    def __init__(self, din, dmid, dout, *, rngs): 
        self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
        self.bn = nnx.BatchNorm(dmid, rngs=rngs)
        self.linear2 = nnx.Linear(dmid, dout, rngs=rngs) 
    def __call__(self, x): 
        x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))
        return self.linear2(x)

# Instantiate model
my_callable = MLP(din=30, dmid=20, dout=10, rngs=nnx.Rngs(0))

# Convert to ONNX
onnx_model = to_onnx(my_callable, [("B", 30)])

# Save the model
onnx.save_model(onnx_model, "my_callable.onnx")

🔎 See it visualized: jax_callable.onnx

---

🧠 ONNX Functions — Minimal Example

ONNX functions help encapsulate reusable subgraphs. Simply use the @onnx_function decorator to make your callable an ONNX function. Just an @onnx_function decorator to make your callable an ONNX function

from onnx import save_model
from flax import nnx
from jax2onnx import onnx_function, to_onnx

# just an @onnx_function decorator to make your callable an ONNX function
@onnx_function
class MLPBlock(nnx.Module):
  def __init__(self, dim, *, rngs):
    self.linear1 = nnx.Linear(dim, dim, rngs=rngs)
    self.linear2 = nnx.Linear(dim, dim, rngs=rngs)
    self.batchnorm = nnx.BatchNorm(dim, rngs=rngs)
  def __call__(self, x):
    return nnx.gelu(self.linear2(self.batchnorm(nnx.gelu(self.linear1(x)))))

# Use it inside another module
class MyModel(nnx.Module):
  def __init__(self, dim, *, rngs):
    self.block1 = MLPBlock(dim, rngs=rngs)
    self.block2 = MLPBlock(dim, rngs=rngs)
  def __call__(self, x):
    return self.block2(self.block1(x))

callable = MyModel(256, rngs=nnx.Rngs(0))
model = to_onnx(callable, [(100, 256)])
save_model(model, "docs/onnx/model_with_function.onnx")

🔎 See it visualized: model_with_function.onnx

---

📅 Roadmap and Releases

Planned Versions

• Ongoing: Expanding JAX component coverage.
• 0.5.0: Some more ONNX function support ... batch dims, function reuse, make graph optimizer work within functions, allow user friendly var names

Current Productive Version

• 0.4.2 (PyPI): Cleanup and fixes to the basic ONNX function release.

Past Versions

• 0.4.1 (ONNX functions): Introducing simple ONNX function support. Making use of ONNX functions is easy for the user: just a @onnx_function decorator making a callable an ONNX function. Each @onnx_function decorator creates a new ONNX function instance on the call graph.
• 0.3.2: relaxed the minimum Python version to 3.10.
• 0.3.0: Streamlined the plugin system with automatic registration and simplified integration of custom primitives.
• 0.2.0 (First PyPI Release): Rebased the implementation on jaxpr, improving usability and adding low-level lax components.
• 0.1.0 (Initial Approach, Not Released to PyPI): Produced ONNX exports for some nnx components and nnx-based examples, including a VisualTransformer.

---

❓ Troubleshooting

If conversion doesn't work out of the box, it could be due to:

• Non-dynamic function references:
  JAXPR-based conversion requires function references to be resolved dynamically at call-time.
  Solution: Wrap your function call inside a lambda to enforce dynamic resolution:

  my_dynamic_callable_function = lambda x: original_function(x)
  

• Unsupported primitives:
  The callable may use a primitive not yet or not fully supported by jax2onnx.
  Solution: Write a plugin to handle the unsupported function (this is straightforward!).

---

🧩 Supported JAX/ONNX Components

JAX Component	ONNX Components	Testcases	Since
add	Add	add ✅	v0.1.0
concat	Concat	concat ✅ concat_abstract_middle_dim_dynamic ✅ concat_abstract_middle_dim ✅	v0.1.0
einsum	Einsum	einsum ✅ einsum_preferred_element_type ✅ einsum_matmul ✅ einsum_dynamic_dynamic ✅ einsum_dynamic ✅ einsum_dynamic_matmul_dynamic ✅ einsum_dynamic_matmul ✅ einsum_transpose ✅ einsum_dynamic_transpose_dynamic ✅ einsum_dynamic_transpose ✅ einsum_dynamic_matmul2_dynamic ✅ einsum_dynamic_matmul2 ✅ einsum_dynamic_matmul3_dynamic ✅ einsum_dynamic_matmul3 ✅ einsum_outer_product ✅ einsum_trace ✅ einsum_sum ✅ einsum_broadcast ✅ einsum_reduce ✅ einsum_permute ✅ einsum_dynamic_outer_dynamic ✅ einsum_dynamic_outer ✅ einsum_dynamic_reduce_dynamic ✅ einsum_dynamic_reduce ✅	v0.1.0
matmul	MatMul	matmul_2d ✅ matmul_1d_2d ✅ matmul_2d_1d ✅ matmul_dynamic_dynamic ✅ matmul_dynamic ✅ matmul_dynamic_a_dynamic ✅ matmul_dynamic_a ✅ matmul_dynamic_b_dynamic ✅ matmul_dynamic_b ✅ matmul_1d ✅ matmul_3d ✅	v0.1.0
reshape	Reshape	reshape_1 ✅ reshape_2 ✅ reshape_3 ✅ reshape_4_dynamic ✅ reshape_4 ✅ reshape_to_scalar ✅ reshape_from_scalar ✅	v0.1.0
shape	Shape	shape_basic ✅ shape_dynamic_dynamic ✅ shape_dynamic ✅	0.4.0
squeeze	Squeeze	squeeze_single_dim ✅ squeeze_multiple_dims ✅ squeeze_vit_output ✅ squeeze_dynamic_batch_dynamic ✅ squeeze_dynamic_batch ✅ squeeze_all_dims ✅ squeeze_negative_axis ✅ squeeze_negative_axis_tuple ✅ squeeze_dynamic_and_negative_axis_dynamic ✅ squeeze_dynamic_and_negative_axis ✅	v0.1.0
tile	Tile	tile_repeats_tensor ✅ tile_a ✅ tile_b ✅ tile_c ✅ tile_d ✅ tile_dynamic_dynamic ✅ tile_dynamic ✅ tile_pad ✅	v0.1.0
transpose	Transpose	transpose_basic ✅ transpose_reverse ✅ transpose_4d ✅ transpose_square_matrix ✅ transpose_high_dim ✅ transpose_no_axes ✅ transpose_dynamic_dynamic ✅ transpose_dynamic ✅	v0.1.0
add	Add	add ✅	v0.2.0
argmax	ArgMax	argmax_test1 ✅ argmax_test2 ✅	v0.2.0
argmin	ArgMin	argmin_test1 ✅ argmin_test2 ✅	v0.2.0
broadcast_in_dim	Expand	broadcast_in_dim ✅ broadcast_in_dim_2d_to_3d ✅ broadcast_in_dim_scalar ✅	v0.2.0
concatenate	Concat	concatenate ✅ concatenate_axis1 ✅ concatenate_dynamic_dynamic ✅ concatenate_dynamic ✅ concatenate_3d ✅	v0.2.0
conv	Conv	conv ✅ conv2 ✅	v0.2.0
convert_element_type	Cast	convert_element_type ✅	v0.2.0
device_put	Identity	device_put_array ✅ device_put_scalar ✅	v0.4.0
div	Div	div ✅	v0.2.0
dot_general	MatMul	dot_general ✅	v0.2.0
dynamic_slice	Slice	dynamic_slice_test1 ✅ dynamic_slice_2d ✅ dynamic_slice_3d ✅	v0.1.0
eq	Equal	eq ✅	v0.2.0
exp	Exp	exp ✅	v0.2.0
gather	Gather	gather ✅	v0.2.0
gt	Greater	gt ✅	v0.2.0
integer_pow	Pow	integer_pow ✅	v0.2.0
log	Log	log ✅	v0.2.0
lt	Less	lt ✅	v0.2.0
max	Max	max ✅	v0.2.0
min	Min	min_test1 ✅	v0.1.0
mul	Mul	mul_test1 ✅ mul_test2 ✅	v0.1.0
ne	Equal Not	ne ✅	v0.2.0
neg	Neg	neg ✅	v0.2.0
reduce_max	ReduceMax	reduce_max ✅	v0.2.0
reduce_min	ReduceMin	reduce_min ✅	v0.2.0
reduce_sum	ReduceSum	reduce_sum ✅	v0.2.0
reshape	Reshape	reshape ✅	v0.2.0
slice	Slice	slice_test1 ✅	v0.1.0
sort	TopK	sort_1d ✅ sort_1d_empty ✅ sort_1d_single ✅ sort_1d_larger ✅ sort_1d_specific_values ✅	v0.2.0
sqrt	Sqrt	sqrt ✅	v0.2.0
square	Mul	square ✅	v0.2.0
squeeze	Squeeze	squeeze ✅	v0.2.0
stop_gradient	Identity	stop_gradient ✅	v0.2.0
sub	Sub	sub_test1 ✅ sub_test2 ✅	v0.1.0
tanh	Tanh	tanh ✅	v0.2.0
transpose	Transpose	transpose_basic ✅	v0.2.0
softmax	Softmax	softmax ✅ softmax_2d ✅ softmax_3d ✅	v0.1.0
avg_pool	AveragePool Transpose	avg_pool ✅ avg_pool_same_padding ✅ avg_pool_default_padding ✅ avg_pool_stride1 ✅ avg_pool_large_window ✅ avg_pool_single_batch ✅ avg_pool_dynamic_batch_dynamic ✅ avg_pool_dynamic_batch ✅ avg_pool_stride_none ✅ avg_pool_count_include_pad_false ✅	v0.1.0
batch_norm	BatchNormalization	batch_norm ✅ batch_norm_2 ✅ batch_norm_3d ✅ batch_norm_float64 ✅ batch_norm_single_batch ✅ batch_norm_2d_train ✅ batch_norm_4d_use_bias ✅ batch_norm_4d_use_scale ✅ batch_norm_momentum ✅ batch_norm_epsilon ✅ batch_norm_float32 ✅ batch_norm_3d_train ✅ batch_norm_single_batch_train ✅	v0.1.0
conv	Conv Transpose	conv_basic_bias_dynamic ✅ conv_basic_bias ✅ conv_basic_bias_2 ✅ conv_basic_bias_3 ✅ conv_stride2_bias ✅ conv_no_bias_dynamic ✅ conv_no_bias ✅ conv_valid_padding ✅ conv_stride1 ✅ conv_stride2 ✅ conv_different_kernel ✅ conv_float64 ✅ conv_single_batch ✅ conv_large_batch ✅	v0.1.0
conv_transpose	ConvTranspose	conv_transpose_valid_padding ✅ conv_transpose_circular_padding ✅	v0.3.0
dot_product_attention	Cast Div Einsum Gather Shape Softmax Sqrt	dpa_basic ✅ dpa_diff_heads_embed ✅ dpa_batch4_seq16 ✅ dpa_float64 ✅ dpa_heads1_embed4 ✅ dpa_heads8_embed8 ✅ dpa_batch1_seq2 ✅ dpa_batch8_seq4 ✅ dpa_axis1 ✅	v0.1.0
dropout	Dropout	dropout_init_params ✅ dropout_call_params ✅	v0.1.0
einsum	Add Einsum	einsum_module_with_bias ✅ einsum_module_no_bias ✅	v0.4.2
elu	Elu	elu ✅	v0.1.0
gelu	Gelu	gelu ✅ gelu_1 ✅ gelu_2 ✅ gelu_3 ✅	v0.1.0
group_norm	GroupNormalization	group_norm ✅ group_norm_2 ✅	v0.3.0
layer_norm	LayerNormalization	layer_norm ✅ layer_norm_multiaxis ✅	v0.1.0
leaky_relu	LeakyRelu	leaky_relu ✅	v0.1.0
linear	Gemm Reshape	linear_2d ✅ linear_dynamic ✅ linear ✅	v0.1.0
linear_general	Gemm Reshape	linear_general_dynamic ✅ linear_general ✅ linear_general_2 ✅ linear_general_3 ✅ linear_general_4 ✅	v0.1.0
log_softmax	LogSoftmax	log_softmax ✅	v0.1.0
max_pool	MaxPool Transpose	max_pool ✅ max_pool_same_padding ✅	v0.1.0
relu	Relu	relu ✅ relu_2 ✅	v0.1.0
rms_norm	RMSNormalization	rms_norm ✅ rms_norm_2 ✅	v0.3.0
sigmoid	Sigmoid	sigmoid ✅	v0.1.0
softmax	Softmax	softmax ✅	v0.1.0
softplus	Softplus	softplus ✅	v0.1.0
tanh	Tanh	tanh ✅	v0.1.0

Legend:
✅ = Passed
❌ = Failed
➖ = No testcase yet

---

🎯 Examples

Component	Description	Testcases	Since
AutoEncoder	A simple autoencoder example.	simple_autoencoder ✅	v0.2.0
CNN	A simple convolutional neural network (CNN).	simple_cnn ✅	v0.1.0
ClassificationHead	Classification head for Vision Transformer	classification_head ✅	v0.4.0
ConcatClsToken	Concatenate CLS token to the input embedding	concat_cls_token ✅	v0.4.0
ConvEmbedding	Convolutional Token Embedding for MNIST with hierarchical downsampling.	mnist_conv_embedding ✅	v0.1.0
FeedForward	MLP in Transformer	feed_forward ✅	v0.1.0
MLP	A simple Multi-Layer Perceptron (MLP) with BatchNorm, Dropout, and GELU activation.	simple_mlp_dynamic ✅ simple_mlp ✅ simple_mlp_with_call_params_dynamic ✅ simple_mlp_with_call_params ✅	v0.1.0
MultiHeadAttention	This is a multi-head attention module implemented by Flax/nnx that has no ONNX correspondent on the same granularity.	multihead_attention ✅	v0.2.0
PatchEmbedding	Cutting the image into patches and linearly embedding them.	patch_embedding ✅	v0.1.0
PositionalEmbedding	Add positional embedding to the input embedding	positional_embedding ✅	v0.4.0
TransformerBlock	Transformer from 'Attention Is All You Need.'	transformer_block ✅	v0.1.0
TransformerStack	Stack of Transformer blocks	transformer_stack ✅	v0.1.0
VisionTransformer	A Vision Transformer (ViT) model for MNIST with configurable embedding type.	vit_conv_embedding ✅ vit_patch_embedding ✅	v0.2.0
onnx_functions_000	one function on an outer layer.	000_one_function_on_outer_layer ✅	v0.4.0
onnx_functions_001	one function on an inner layer.	001_one_function_inner ✅	v0.4.0
onnx_functions_002	two nested functions.	002_two_nested_functions ✅	v0.4.0
onnx_functions_003	two nested functions.	003_two_simple_nested_functions ✅	v0.4.0
onnx_functions_004	nested function plus component	004_nested_function_plus_component ✅	v0.4.0
onnx_functions_005	nested function plus more components	005_nested_function_plus_component ✅	v0.4.0
onnx_functions_006	one function on an outer layer.	006_one_function_outer ✅	v0.4.0
onnx_functions_008	transformer block with nested mlp block no call parameter	008_transformer_block ✅	v0.4.0
onnx_functions_009	transformer block using decorator on class and function	009_transformer_block ✅	v0.4.0
onnx_functions_010	transformer stack	010_transformer_stack ✅	v0.4.0
onnx_functions_012	Vision Transformer (ViT)	012_vit_conv_embedding ✅	v0.4.0
onnx_functions_013	Vision Transformer (ViT)	013_vit_conv_embedding_with_call_params ✅ 013_vit_conv_embedding_with_internal_call_params ✅	v0.4.0
onnx_functions_014	one function on an outer layer.	014_one_function_with_input_param_with_default_value ✅ 014_one_function_without_input_param_with_default_value ✅	v0.4.0
onnx_functions_015	one function on an outer layer.	015_one_function_with_input_param_without_default_value ✅	v0.4.0
onnx_functions_016	nested function plus more components	016_internal_function_with_input_param_with_default_value ✅	v0.4.0

---

📌 Dependencies

Versions of Major Dependencies:

Library	Versions
JAX	0.5.3
Flax	0.10.5
onnx	1.17.0
onnxruntime	1.21.0

Note: For more details, check pyproject.toml.

---

⚠️ Limitations

• Currently not all JAX/Flax components are supported (you can easily help expand this coverage!).
• Function references need dynamic resolution at call-time.
• ONNX graph composition is done in-memory before saving to disk, potentially causing memory issues with very large models.

---

🤝 How to Contribute

We warmly welcome contributions!

How you can help:

• Add a plugin: Extend jax2onnx by writing a simple Python file in jax2onnx/plugins. a custom primitive or an example.
• Bug fixes & improvements: PRs and issues are always welcome.

---

💾 Installation

Install from PyPI:

pip install jax2onnx  

Or get the latest development version from TestPyPI:

pip install -i https://test.pypi.org/simple/ jax2onnx

---

📜 License

This project is licensed under the Apache License, Version 2.0. See LICENSE for details.

---

🌟 Special Thanks

Special thanks to @lutzroeder for making shapes internal to ONNX function visible in his great Netron viewer.

• ONNX: Function value_info support #1447

Special thanks to the community members involved in:

• Flax Feature Request #4430
• JAX Feature Request #26430

A huge thanks especially to @limarta, whose elegant jaxpr-to-ONNX demonstration significantly inspired this project.

---

Happy converting! 🎉