jax2onnx
export JAX to ONNX
File Explorer
Download Latest Version (.zip)- ci.yml
- docs.yml
- nightly-qs.yml
- dependabot.yml
- acknowledgements.md
- contributing.md
- dependencies.md
- license.md
- past_versions.md
- roadmap.md
- complex_numbers.md
- expect_graph_reference.md
- ir_optimizer.md
- ir_reflection_guidelines.md
- onnx_ir_builder.md
- onnx_ir_typing_how_to.md
- subgraph_input_handling.md
- typing_improvements.md
- artifact_publishing.md
- coverage_generation.md
- sota_example_maintenance.md
- architecture.md
- plugin_system.md
- jax2onnx.svg
- extra.css
- compare_meta_vs_jax2onnx.md
- getting_weights.md
- nnx_exports.md
- getting_weights.md
- maxdiffusion.md
- maxtext.md
- api.md
- browser_wasm.md
- equinox_nn_coverage.md
- examples.md
- flax_api_coverage.md
- getting_started.md
- jax_lax_coverage.md
- jax_numpy_coverage.md
- known_limitations.md
- layout_optimization.md
- onnx_functions.md
- onnx_operator_coverage.md
- rl_policy_deployment.md
- stacktrace_metadata.md
- supported_components.md
- validation.md
- index.md
- converter_design_review.md
- deployment_readiness_roadmap.md
- __init__.py
- jax.py
- __init__.py
- conversion_api.py
- function_scope.py
- ir_builder.py
- ir_constants.py
- ir_context.py
- ir_optimizations.py
- ir_postprocess.py
- lower_dimexpr.py
- lowering_dispatch.py
- optimizer_graph_utils.py
- output_binding.py
- typing_support.py
- __init__.py
- depth_to_space.py
- space_to_depth.py
- __init__.py
- adaptive_pool.py
- avg_pool.py
- batch_norm.py
- conv.py
- dropout.py
- embedding.py
- group_norm.py
- identity.py
- lambda.py
- layer_norm.py
- linear.py
- max_pool.py
- multihead_attention.py
- pool.py
- prelu.py
- recurrent.py
- rms_norm.py
- rotary_positional_embedding.py
- sequential.py
- spectral_norm.py
- weight_norm.py
- __init__.py
- __init__.py
- control_primitives.py
- __init__.py
- __init__.py
- dino.py
- gpt_oss.py
- mlp.py
- simple_linear.py
- __init__.py
- reconstructed_patterns.py
- __init__.py
- fori_loop.py
- issue18.py
- select.py
- sort.py
- __init__.py
- cond_scatter_add_mul.py
- cond_scatter_repro.py
- remat2.py
- scatter_window.py
- two_times_silu.py
- __init__.py
- cnn.py
- mlp.py
- mlp_sequential.py
- __init__.py
- maxdiffusion_models.py
- __init__.py
- maxtext_models.py
- __init__.py
- autoencoder.py
- channel_attention_resblock.py
- cnn.py
- depth_to_space.py
- dinov3.py
- exclusive_self_attention.py
- fori_loop.py
- gpt.py
- gpt_oss_flax.py
- gru_cell.py
- mlp.py
- multi_head_attention.py
- nested_residual_group.py
- sequential.py
- simple_clip.py
- transformer_decoder_with_sequential.py
- transformer_decoder_without_sequential.py
- vit.py
- vit_flat.py
- __init__.py
- onnx_functions_000.py
- onnx_functions_001.py
- onnx_functions_002.py
- onnx_functions_003.py
- onnx_functions_004.py
- onnx_functions_005.py
- onnx_functions_006.py
- onnx_functions_007.py
- onnx_functions_008.py
- onnx_functions_009.py
- onnx_functions_010.py
- onnx_functions_012.py
- onnx_functions_013.py
- onnx_functions_014.py
- onnx_functions_015.py
- onnx_functions_016.py
- onnx_functions_017.py
- __init__.py
- policies.py
- __init__.py
- __init__.py
- activation.py
- avg_pool.py
- batch_norm.py
- conv.py
- conv_local.py
- conv_transpose.py
- dense.py
- dense_general.py
- dot_product_attention.py
- dropout.py
- einsum.py
- embed.py
- group_norm.py
- instance_norm.py
- layer_norm.py
- max_pool.py
- min_pool.py
- multi_head_attention.py
- multi_head_dot_product_attention.py
- pool.py
- prelu.py
- recurrent.py
- rms_norm.py
- self_attention.py
- spectral_norm.py
- weight_norm.py
- __init__.py
- avg_pool.py
- batch_norm.py
- combine_masks.py
- conv.py
- dot_product_attention.py
- dropout.py
- einsum.py
- elu.py
- embed.py
- flip_sequences.py
- gelu.py
- glu.py
- group_norm.py
- hard_tanh.py
- layer_norm.py
- leaky_relu.py
- linear.py
- linear_general.py
- log_sigmoid.py
- log_softmax.py
- lora.py
- max_pool.py
- prelu.py
- relu.py
- rms_norm.py
- sigmoid.py
- silu.py
- softmax.py
- softplus.py
- tanh.py
- __init__.py
- test_utils.py
- __init__.py
- custom_jvp_call.py
- custom_vjp_call.py
- dim_as_value.py
- jit.py
- name.py
- __init__.py
- _common.py
- resize.py
- __init__.py
- _arg_utils.py
- _control_flow_utils.py
- _cum_extrema.py
- _index_utils.py
- _opset_utils.py
- _reduce_utils.py
- _reduction_registry.py
- _sort_utils.py
- abs.py
- acos.py
- acosh.py
- add.py
- add_any.py
- and.py
- approx_top_k.py
- argmax.py
- argmin.py
- asin.py
- asinh.py
- atan.py
- atan2.py
- atanh.py
- bessel_i0e.py
- bessel_i1e.py
- betainc.py
- bitcast_convert_type.py
- bitwise_not.py
- broadcast_in_dim.py
- cbrt.py
- ceil.py
- cholesky.py
- cholesky_update.py
- clamp.py
- clz.py
- complex.py
- concatenate.py
- cond.py
- conj.py
- conv.py
- convert_element_type.py
- copy.py
- cos.py
- cosh.py
- cumlogsumexp.py
- cummax.py
- cummin.py
- cumprod.py
- cumsum.py
- custom_linear_solve.py
- device_put.py
- digamma.py
- div.py
- dot_general.py
- dynamic_slice.py
- dynamic_update_slice.py
- eig.py
- eigh.py
- empty.py
- eq.py
- erf.py
- erf_inv.py
- erfc.py
- exp.py
- exp2.py
- expm1.py
- fft.py
- floor.py
- fori_loop.py
- gather.py
- gather_compile.py
- gather_helpers.py
- ge.py
- gt.py
- hessenberg.py
- householder_product.py
- igamma.py
- imag.py
- integer_pow.py
- iota.py
- is_finite.py
- le.py
- lgamma.py
- log.py
- log1p.py
- logistic.py
- lt.py
- lu.py
- lu_pivots_to_permutation.py
- max.py
- min.py
- mul.py
- ne.py
- neg.py
- nextafter.py
- optimization_barrier.py
- or.py
- ormqr.py
- pad.py
- pjit.py
- polygamma.py
- population_count.py
- pow.py
- qr.py
- ragged_dot_general.py
- real.py
- reduce.py
- reduce_and.py
- reduce_max.py
- reduce_min.py
- reduce_or.py
- reduce_precision.py
- reduce_prod.py
- reduce_sum.py
- reduce_window.py
- reduce_window_max.py
- reduce_window_sum.py
- reduce_xor.py
- rem.py
- remat2.py
- reshape.py
- rev.py
- rng_bit_generator.py
- rng_uniform.py
- round.py
- rsqrt.py
- scan.py
- scatter.py
- scatter_add.py
- scatter_max.py
- scatter_min.py
- scatter_mul.py
- scatter_sub.py
- scatter_utils.py
- schur.py
- select.py
- select_n.py
- shard_map.py
- sharding_constraint.py
- shift_left.py
- shift_right_arithmetic.py
- shift_right_logical.py
- sign.py
- sin.py
- sinh.py
- slice.py
- sort.py
- split.py
- sqrt.py
- square.py
- squeeze.py
- stop_gradient.py
- sub.py
- svd.py
- symmetric_product.py
- tan.py
- tanh.py
- top_k.py
- transpose.py
- triangular_solve.py
- tridiagonal.py
- tridiagonal_solve.py
- while_loop.py
- xor.py
- zeta.py
- __init__.py
- truncated_normal.py
- __init__.py
- _builder_utils.py
- celu.py
- dot_product_attention.py
- elu.py
- gelu.py
- glu.py
- hard_sigmoid.py
- hard_swish.py
- hard_tanh.py
- hardmax.py
- identity.py
- leaky_relu.py
- log1mexp.py
- log_sigmoid.py
- log_softmax.py
- logmeanexp.py
- logsumexp.py
- mish.py
- one_hot.py
- relu.py
- relu6.py
- scaled_dot_general.py
- scaled_matmul.py
- selu.py
- sigmoid.py
- silu.py
- softmax.py
- softplus.py
- softsign.py
- sparse_plus.py
- sparse_sigmoid.py
- squareplus.py
- standardize.py
- tanh.py
- thresholded_relu.py
- __init__.py
- _common.py
- _reduction_utils.py
- _unary_utils.py
- abs.py
- acos.py
- acosh.py
- add.py
- all.py
- amax.py
- amin.py
- any.py
- arange.py
- argmax.py
- argmin.py
- asin.py
- asinh.py
- atan.py
- atan2.py
- atanh.py
- bitwise_and.py
- bitwise_left_shift.py
- bitwise_not.py
- bitwise_or.py
- bitwise_right_shift.py
- bitwise_xor.py
- ceil.py
- clip.py
- composite_metadata.py
- composite_metadata_batch2.py
- composite_metadata_batch3.py
- composite_metadata_batch4.py
- composite_metadata_batch5.py
- composite_metadata_batch6.py
- compress.py
- concatenate.py
- conj.py
- copysign.py
- cos.py
- cosh.py
- cumprod.py
- cumsum.py
- diag.py
- diagonal.py
- digitize.py
- divide.py
- dot.py
- einsum.py
- equal.py
- exp.py
- exp2.py
- expm1.py
- eye.py
- fabs.py
- fft.py
- floor.py
- floor_divide.py
- fmod.py
- frexp.py
- full.py
- greater.py
- greater_equal.py
- histogram.py
- histogram2d.py
- histogramdd.py
- interp.py
- invert.py
- isfinite.py
- ldexp.py
- left_shift.py
- less.py
- less_equal.py
- linalg_det.py
- linalg_inv.py
- linalg_norm.py
- linalg_solve.py
- linalg_tensorinv.py
- linalg_tensorsolve.py
- linspace.py
- log.py
- matmul.py
- max.py
- maximum.py
- mean.py
- min.py
- minimum.py
- moveaxis.py
- nancumprod.py
- ones.py
- outer.py
- pad.py
- polyfit.py
- pow.py
- prod.py
- reshape.py
- right_shift.py
- roots.py
- searchsorted.py
- select.py
- shape.py
- sign.py
- sin.py
- sinh.py
- size.py
- sort.py
- spacing.py
- split.py
- sqrt.py
- squeeze.py
- stack.py
- sum.py
- take.py
- tan.py
- tanh.py
- tile.py
- transpose.py
- trilu.py
- unique.py
- unstack.py
- where.py
- windows.py
- zeros.py
- __init__.py
- bernoulli.py
- categorical.py
- normal.py
- random_bits.py
- random_fold_in.py
- random_seed.py
- __init__.py
- _autodiff_utils.py
- _batching_compat.py
- _batching_utils.py
- __init__.py
- _attention_utils.py
- _axis0_utils.py
- _complex_utils.py
- _ir_shapes.py
- _loop_extent_meta.py
- _patching.py
- _post_check_onnx_graph.py
- _utils.py
- plugin_system.py
- __init__.py
- issue139.py
- __init__.py
- issue_145.py
- __init__.py
- create_checkpoint_example.py
- export_onnx.py
- __init__.py
- export_onnx.py
- __init__.py
- repro.py
- __init__.py
- issue52_broadcast_bug.py
- issue52_feedforward_payload.npz
- issue52_loop_concat_bug.py
- issue52_scatter_payload_repro.py
- issue52_stack_bug.py
- __init__.py
- dino_01.py
- example_1.py
- example_2.py
- issue_206.py
- debug.py
- parameter_validation.py
- shape_poly.py
- __init__.py
- _deployment_report.py
- ir_utils.py
- quickstart.py
- quickstart_functions.py
- user_interface.py
- _coverage_generation.py
- audit_ir_dynamic_access.py
- check_forbidden_reflection.py
- check_ir_builder_usage.py
- check_typing.sh
- check_variable_annotations.py
- compare_dinov3_embeddings.py
- compare_eqx_onnx_gpt_oss.py
- compare_meta_vs_jax2onnx.py
- convert_dinov3_from_equimo.py
- convert_reduction_tests.py
- debug_attn_parity.py
- debug_norm_parity.py
- debug_rope_parity.py
- dino_block_deltas.py
- dino_block_token_deltas.py
- emit_expect_graph.py
- ensure_path_markers.py
- export_dinov3_pretrained.py
- export_eqx_dino_example_with_mapped_weights.py
- export_eqx_gpt_oss_example_with_mapped_weights.py
- export_eqx_gpt_oss_subset.py
- export_eqx_gpt_oss_to_onnx.py
- export_flax_gpt_oss_params.py
- export_flax_gpt_oss_to_onnx.py
- gen_arch_diagrams.py
- generate_dinov3_reference.py
- generate_eqx_nn_coverage.py
- generate_flax_api_coverage.py
- generate_jnp_operator_coverage.py
- generate_lax_operator_coverage.py
- generate_onnx_operator_coverage.py
- generate_readme.py
- generate_readme.sh
- generate_tests.py
- gpt_oss_routing_parity.py
- install_minimum_dependencies.py
- map_equimo_dino_weights.py
- onnxruntime_web_validation_utils.mjs
- probe_eqx_gpt_oss_parity.py
- probe_eqx_gpt_oss_parity_real.py
- probe_flax_gpt_oss_parity.py
- report_rng_traces.py
- run_all_checks.sh
- run_flax_gpt_oss_onnx.py
- run_mypy.py
- run_onnx_only.py
- run_onnxruntime_web_chrome_full.sh
- run_onnxruntime_web_chrome_smoke.sh
- run_onnxruntime_web_full.sh
- run_onnxruntime_web_smoke.sh
- sync_onnx_models.sh
- validate_onnxruntime_web.mjs
- validate_onnxruntime_web_chrome.mjs
- version.py
- favicon.png
- lunr.ar.min.js
- lunr.da.min.js
- lunr.de.min.js
- lunr.du.min.js
- lunr.el.min.js
- lunr.es.min.js
- lunr.fi.min.js
- lunr.fr.min.js
- lunr.he.min.js
- lunr.hi.min.js
- lunr.hu.min.js
- lunr.hy.min.js
- lunr.it.min.js
- lunr.ja.min.js
- lunr.jp.min.js
- lunr.kn.min.js
- lunr.ko.min.js
- lunr.multi.min.js
- lunr.nl.min.js
- lunr.no.min.js
- lunr.pt.min.js
- lunr.ro.min.js
- lunr.ru.min.js
- lunr.sa.min.js
- lunr.stemmer.support.min.js
- lunr.sv.min.js
- lunr.ta.min.js
- lunr.te.min.js
- lunr.th.min.js
- lunr.tr.min.js
- lunr.vi.min.js
- lunr.zh.min.js
- tinyseg.js
- wordcut.js
- search.2c215733.min.js
- search.2c215733.min.js.map
- bundle.79ae519e.min.js
- bundle.79ae519e.min.js.map
- main.484c7ddc.min.css
- main.484c7ddc.min.css.map
- palette.ab4e12ef.min.css
- palette.ab4e12ef.min.css.map
- _mkdocstrings.css
- search_index.json
- 404.html
- index.html
- objects.inv
- sitemap.xml
- sitemap.xml.gz
- __init__.py
- test_broadcast_and_reshape_shared_shape_ssa_regression.py
- test_broadcast_in_dim_in_loop_ssa.py
- test_broadcast_in_dim_loop_ssa_regression.py
- test_broadcast_in_dim_shape_helper_metadata.py
- test_broadcast_in_dim_shapeof_ssa_regression_strict.py
- test_broadcast_in_dim_ssa_unique.py
- test_broadcast_in_dim_two_adds_shapeof_ssa_regression.py
- __init__.py
- test_complex_utils.py
- test_context_typing_compat.py
- test_conversion_context_helpers.py
- test_conversion_dtype.py
- test_conversion_finalization_helpers.py
- test_conversion_guardrails.py
- test_conversion_lowering_helpers.py
- test_conversion_output_helpers.py
- test_conversion_primitive_recording.py
- test_conversion_trace_helpers.py
- test_default_opset.py
- test_duplicate_params.py
- test_dynamic_batch.py
- test_exact_dtype_casts.py
- test_f64_constant_promotion.py
- test_initializer_shape_guard.py
- test_io_names.py
- test_ir_attr_conversion.py
- test_ir_builder.py
- test_ir_clone.py
- test_ir_constants.py
- test_ir_context_nodes.py
- test_ir_postprocess.py
- test_jaxpr_converter_interaction_with_builder.py
- test_layout_optimization.py
- test_lowering_dispatch.py
- test_nested_params.py
- test_no_null_node.py
- test_onnx_function_call_params.py
- test_onnx_function_deterministic_param.py
- test_onnx_function_unique.py
- test_output_binding.py
- test_remove_redundant_reshapes.py
- test_subgraph_inputs.py
- test_symbolic_dim_name.py
- test_value_metadata.py
- test_eqx_dino_rotary_lowering.py
- test_io_names_examples.py
- test_nnx_dino_parity.py
- __init__.py
- test_ad_backfill_idempotence.py
- test_ad_registration_helpers_only.py
- test_ad_rule_completeness.py
- test_ad_rule_forwarding.py
- test_ad_rule_no_override.py
- test_coverage_docs.py
- test_coverage_generators.py
- test_deployment_readiness_report.py
- test_do_not_skip_numeric_validation.py
- test_generate_tests_entrypoint.py
- test_ir_builder_contracts.py
- test_ir_optimizations.py
- test_ir_roundtrip.py
- test_ir_utils.py
- test_jax_compat.py
- test_jnp_composite_api_reuse.py
- test_minimum_dependency_installer.py
- test_no_onnx_in_converter_plugins.py
- test_onnx_ir_surface.py
- test_opset_sensitive_reduction_axes.py
- test_optional_ormqr_discovery.py
- test_post_check_onnx_graph.py
- test_reduce_ops_coverage.py
- test_since_equal_in_convert.py
- __init__.py
- issue52_broadcast_fixture.py
- issue52_loop_concat_fixture.py
- issue52_scatter_payload_fixture.py
- __init__.py
- test_loop_add_shapeof_displayname_collision.py
- test_loop_add_shapeof_displayname_collision_dev.py
- test_loop_body_io_relaxation.py
- test_loop_body_loosen_shapes_regression.py
- test_loop_body_shapeof_ssa_regression_strict.py
- test_loop_body_ssa_dedup.py
- test_loop_broadcast_extent_regression.py
- test_loop_concat_extent_regression.py
- test_loop_consts_and_subset.py
- test_loop_ff_like_broadcast_mul_inference_regression.py
- test_loop_ff_like_broadcast_mul_regression.py
- test_loop_gather_dtype_regression.py
- test_loop_gathernd_nested_dtype_regression.py
- test_loop_mul_shape_inference_regression.py
- test_loop_output_arity.py
- test_loop_same_rank_broadcast_mul_regression.py
- test_loop_scatter_payload_extent_regression.py
- test_loop_scatter_payload_regression.py
- test_loop_shape_helper_uniqueness_regression.py
- test_loop_shape_inference_regression.py
- test_loop_shapeof_ssa_regression.py
- test_loop_shapeof_ssa_regression_stricter.py
- test_nested_loop_loosen_shapes_regression.py
- test_nested_loop_mul_pad_slice.py
- test_vmapped_while_masking.py
- test_while_control_flow_regressions.py
- __init__.py
- test_pad_initializer_in_loop.py
- issue52_feedforward_payload.npz
- __init__.py
- test_scan_body_loosen_shapes_regression.py
- test_scan_dynamic_len.py
- test_scan_loop_no_xs_regression.py
- test_scan_loop_output_arity_pruned.py
- test_scan_shapeof_ssa_regression.py
- test_scan_shared_scalar_two_lengths.py
- test_scan_tripaxis_dynamic.py
- test_scan_two_tripcounts_dynamic.py
- __init__.py
- test_padding_and_expected_shape.py
- test_where_guardrails.py
- __init__.py
- test_dynamic_update_slice_nd.py
- __init__.py
- test_no_duplicate_shape_helpers_anywhere.py
- __init__.py
- capability_matrix.py
- test_bfloat16_capability_matrix.py
- test_concat_mixed_int_float_and_take.py
- test_cumsum_autodiff.py
- test_equinox_interop_regressions.py
- test_feed_forward_like_nested_loops_dtype_regression.py
- test_gather_modes.py
- test_group_norm_stability.py
- test_issue_139.py
- test_issue_173_nested_residual_groups.py
- test_issue_174_dm_pix_depth_to_space.py
- test_issue_203_linear_transpose_add.py
- test_issue_203_linear_transpose_allowlist_ops.py
- test_issue_203_linear_transpose_transpose.py
- test_issue_210_reflect_padding_dynamic.py
- test_issue_212_dynamic_reverse.py
- test_issue_215_lax_split_single_output.py
- test_issue_221_bfloat16_pool.py
- test_issue_223_nnx_silu_swish.py
- test_issue_57_scatter_add_broadcast_window.py
- test_issue_89.py
- test_issue_95_gather_indexing.py
- test_lax_cond_repro.py
- test_masked_gather_trig.py
- test_normalization_export_policy.py
- test_opset23_native_attention_rotary.py
- test_opset27_range_and_resize.py
- test_quickstart.py
- test_return_modes.py
- test_scatter_broadcast_isolation.py
- test_split_batching_and_empty_outputs.py
- test_subgraph_output_leaking.py
- test_vmap_batching_regression.py
- __init__.py
- _initializer_guard.py
- conftest.py
- t_generator.py
- test_gpt_oss_components.py
- .gitignore
- .pre-commit-config.yaml
- .python-version
- AGENTS.md
- LICENSE
- logging_config.py
- mkdocs.yml
- package-lock.json
- package.json
- poetry.lock
- pyproject.toml
- README.md
- uv.lock
// repository documentation
Was this content helpful?
(0 ratings)