Metadata-Version: 2.3
Name: jax2onnx
Version: 0.7.2
Summary: export JAX to ONNX - focus on flax nnx
Author: enpasos
Author-email: matthias.unverzagt@enpasos.ai
Requires-Python: >=3.10
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Dist: einops (>=0.8.1,<0.9.0)
Requires-Dist: equinox (>=0.12.2,<0.13.0)
Requires-Dist: flax (>=0.10.6)
Requires-Dist: jax (>=0.6.0,<0.7.0)
Requires-Dist: ml_dtypes (==0.5.1)
Requires-Dist: netron (>=8.2.9)
Requires-Dist: onnx (>=1.17.0)
Requires-Dist: onnxruntime (>=1.22.0)
Requires-Dist: optax (>=0.2.4)
Requires-Dist: orbax-checkpoint (>=0.11.6)
Requires-Dist: orbax-export (>=0.0.6)
Description-Content-Type: text/markdown

# jax2onnx 🌟

`jax2onnx` converts your JAX/Flax(nnx) functions directly into the ONNX format.


![img.png](https://enpasos.github.io/jax2onnx/images/jax2onnx.png)

## ✨ Key Features

- **Simple API**  
  Easily convert JAX callables—including Flax (NNX) models—into ONNX format using `to_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:

```python
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:  [`my_callable.onnx`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/my_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

```python
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`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/model_with_function.onnx)

---

## 📅 Roadmap and Releases

### **Planned Versions**

* **Ongoing**

  * Expanding coverage of JAX and Flax (NNX) components
  * Enhancing support for **physics-based simulations**
 
* **Upcoming**

  * Advanced ONNX function support, including function reuse, optimized internal graph structure, and improved variable naming for clarity and readability
  * Support for `equinox` models 
  * Integrating `onnx-ir` as a backend to improve ONNX model construction, memory efficiency, and performance
 

### **Current Productive Version**
* **0.7.2** *(PyPI)*:
  * Added support for new primitives: `jnp.split`, `lax.split`, `lax.logistic`.
  * Added examples for `nnx.GRUCell`
  * Improved support for `lax.scatter`, and `lax.while_loop`.

### **Past Versions**
* **0.7.1**:
  * Fixed a bug in the test system ensuring numeric equivalence with Jax.
  * Added support for new primitives: `core.custom_jvp_generic`, `eqx.identity`, `jnp.select`, `jnp.stack`, `jnp.unstack`, `lax.select`, `nn.identiy`, `nn.celu`, `nn.elu`, `nn.gelu`, `nn.relu`, `nn.leaky_relu`, `nn.mish`, `nn.selu`, `nn.sigmoid`, `nn.soft-sign`, `nn.softmax`, `nn.truncated_normal`
* **0.7.0**: 
  * Added a GPT-2 model example based on [nanoGPT](https://github.com/karpathy/nanoGPT/blob/master/model.py), featuring ONNX function support and attention masking
  * New support for `jnp.concatenate`, `jnp.take`, `nnx.Embed`
  * ONNX models are now hosted on Hugging Face 
- **0.6.5**: Improved support for `nnx.batch_norm`, `nnx.group_norm`, `nnx.layer_norm`, `nnx.rms_norm`, `lax.broadcast_in_dim`, `lax.cond`, `lax.fori_loop`,
  `lax.integer_pow`, `lax.scan`, `lax.scatter`, `lax.scatter_add`, `lax.scatter_mul` and `lax.while_loop`;
  added support for `lax.and`, `lax.rem` and `lax.remat2`.
- **0.6.4**: Improved support for `lax.scatter_mul`.
- **0.6.3**: Double precision fixes for  `lax.fori_loop` and `lax.while_loop`. Fixed bugs in `lax.scan` and `jnp.where`.
- **0.6.2**: Fixed bugs in `nnx.conv` and `lax.reshape`; added new primitive `jnp.prod`.
- **0.6.1**: Improved support for `lax.cond` and `lax.select_n`; added new primitives (`lax.reduce_and`, `lax.reduce_or`, `lax.reduce_prod`, `lax.reduce_xor`); and introduced new examples for `jnp.select` and `jnp.sort`. 
- **0.6.0**: Introduced the `enable_double_precision` parameter (default: `False`) to support **physics simulations**, and enhanced handling of `lax.scatter`.
- **0.5.2**: Add support for additional primitives: `jnp.where`, `jnp.arange`, `jnp.linspace`.
- **0.5.1**: Add support for subgraph using primitives: `lax.while_loop`, `lax.cond`, `lax.fori_loop`, `lax.scan`.
- **0.5.0**: Improved dynamic batch dimension handling by leveraging [shape polymorphism](https://docs.jax.dev/en/latest/export/shape_poly.html) for more robust and flexible model export.
  Added support for `jnp.sign`, `jnp.abs`, `jnp.iota` primitives.
- **0.4.4**: Added support for `lax.cos`, `lax.cosh`, `lax.sin`, `lax.sinh` and `lax.scatter` primitives.
- **0.4.3**: Fixed a bug in the validation of JAX callable outputs against their ONNX counterparts. This fix exposed previously hidden failing tests, which are now fixed.
- **0.4.2**: Cleanup and fixes to the basic ONNX function release.
- **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:
  ```python
  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](#how-to-contribute) to handle the unsupported function (this is straightforward!).

---

## 🧩 Supported JAX/ONNX Components


<!-- AUTOGENERATED TABLE START -->

| JAX Component | ONNX Components | Testcases | Since |
|:-------------|:---------------|:---------|:------|
| [core.custom_jvp_generic](Generic passthrough for custom JVP calls) | [CustomJvp]() | [`custom_jvp_square`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/core/custom_jvp_square.onnx) ✅<br>[`custom_jvp_square_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/core/custom_jvp_square_f64.onnx) ✅ | v0.7.1 |
| [core.dim_as_value](https://github.com/jax-ml/jax/blob/main/jax/_src/export/shape_poly.py) | [Cast](https://onnx.ai/onnx/operators/onnx__Cast.html)<br>[Gather](https://onnx.ai/onnx/operators/onnx__Gather.html)<br>[Reshape](https://onnx.ai/onnx/operators/onnx__Reshape.html)<br>[Shape](https://onnx.ai/onnx/operators/onnx__Shape.html) | [`dim_as_value_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/core/dim_as_value_dynamic.onnx) ✅<br>[`dim_as_value`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/core/dim_as_value.onnx) ✅ | v0.5.0 |
| [eqx.identity](https://docs.kidger.site/equinox/api/nn/linear/#equinox.nn.Identity) | [Identity](https://onnx.ai/onnx/operators/onnx__Identity.html) | [`eqx_identity_static`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/eqx/eqx_identity_static.onnx) ✅<br>[`eqx_identity_static_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/eqx/eqx_identity_static_f64.onnx) ✅<br>[`eqx_identity_symbolic_batch_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/eqx/eqx_identity_symbolic_batch_dynamic.onnx) ✅<br>[`eqx_identity_symbolic_batch_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/eqx/eqx_identity_symbolic_batch_dynamic_f64.onnx) ✅<br>[`eqx_identity_symbolic_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/eqx/eqx_identity_symbolic_batch.onnx) ✅<br>[`eqx_identity_symbolic_batch_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/eqx/eqx_identity_symbolic_batch_f64.onnx) ✅ | v0.7.1 |
| [jnp.add](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.add.html) | [Add](https://onnx.ai/onnx/operators/onnx__Add.html) | [`add`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/add.onnx) ✅<br>[`add_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/add_f64.onnx) ✅ | v0.1.0 |
| [jnp.arange](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.arange.html) | [Range](https://onnx.ai/onnx/operators/onnx__Range.html) | [`arange_data_dependent_indices`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_data_dependent_indices.onnx) ✅<br>[`arange_stop_only_concrete_input_val`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_stop_only_concrete_input_val.onnx) ✅<br>[`arange_stop_only_concrete_input_val_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_stop_only_concrete_input_val_f64.onnx) ✅<br>[`arange_start_stop_concrete_input_val`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_start_stop_concrete_input_val.onnx) ✅<br>[`arange_start_stop_concrete_input_val_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_start_stop_concrete_input_val_f64.onnx) ✅<br>[`arange_start_stop_step_concrete_input_val`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_start_stop_step_concrete_input_val.onnx) ✅<br>[`arange_start_stop_step_concrete_input_val_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_start_stop_step_concrete_input_val_f64.onnx) ✅<br>[`arange_float_concrete_input_val`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_float_concrete_input_val.onnx) ✅<br>[`arange_float_concrete_input_val_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_float_concrete_input_val_f64.onnx) ✅<br>[`arange_static_stop_only_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_stop_only_int.onnx) ✅<br>[`arange_static_stop_only_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_stop_only_int_f64.onnx) ✅<br>[`arange_static_stop_only_float`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_stop_only_float.onnx) ✅<br>[`arange_static_stop_only_float_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_stop_only_float_f64.onnx) ✅<br>[`arange_static_start_stop_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_start_stop_int.onnx) ✅<br>[`arange_static_start_stop_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_start_stop_int_f64.onnx) ✅<br>[`arange_static_start_stop_step_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_start_stop_step_int.onnx) ✅<br>[`arange_static_start_stop_step_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_start_stop_step_int_f64.onnx) ✅<br>[`arange_static_empty_result_pos_step`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_empty_result_pos_step.onnx) ✅<br>[`arange_static_empty_result_pos_step_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_empty_result_pos_step_f64.onnx) ✅<br>[`arange_static_empty_result_neg_step`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_empty_result_neg_step.onnx) ✅<br>[`arange_static_empty_result_neg_step_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_empty_result_neg_step_f64.onnx) ✅<br>[`arange_static_negative_step`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_negative_step.onnx) ✅<br>[`arange_static_negative_step_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_negative_step_f64.onnx) ✅<br>[`arange_static_float_step_explicit_dtype`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_float_step_explicit_dtype.onnx) ✅<br>[`arange_static_float_step_explicit_dtype_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_float_step_explicit_dtype_f64.onnx) ✅<br>[`arange_static_float_step_inferred_dtype`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_float_step_inferred_dtype.onnx) ✅<br>[`arange_static_float_step_inferred_dtype_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_float_step_inferred_dtype_f64.onnx) ✅<br>[`arange_static_stop_zero`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_stop_zero.onnx) ✅<br>[`arange_static_stop_zero_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_stop_zero_f64.onnx) ✅<br>[`arange_static_start_equals_stop`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_start_equals_stop.onnx) ✅<br>[`arange_static_start_equals_stop_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_start_equals_stop_f64.onnx) ✅<br>[`arange_static_large_numbers_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_large_numbers_int.onnx) ✅<br>[`arange_static_large_numbers_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/arange_static_large_numbers_int_f64.onnx) ✅ | v0.5.2 |
| [jnp.concatenate](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.concatenate.html) | [Concat](https://onnx.ai/onnx/operators/onnx__Concat.html) | [`concatenate`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate.onnx) ✅<br>[`concatenate_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_f64.onnx) ✅<br>[`concatenate_abstract_middle_dim_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_abstract_middle_dim_dynamic.onnx) ✅<br>[`concatenate_abstract_middle_dim_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_abstract_middle_dim_dynamic_f64.onnx) ✅<br>[`concatenate_abstract_middle_dim`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_abstract_middle_dim.onnx) ✅<br>[`concatenate_abstract_middle_dim_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_abstract_middle_dim_f64.onnx) ✅<br>[`concatenate_tile_and_symbolic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_tile_and_symbolic_dynamic.onnx) ✅<br>[`concatenate_tile_and_symbolic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_tile_and_symbolic_dynamic_f64.onnx) ✅<br>[`concatenate_tile_and_symbolic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_tile_and_symbolic.onnx) ✅<br>[`concatenate_tile_and_symbolic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/concatenate_tile_and_symbolic_f64.onnx) ✅ | v0.1.0 |
| [jnp.einsum](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.einsum.html) | [Einsum](https://onnx.ai/onnx/operators/onnx__Einsum.html) | [`einsum_vector_dot`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_vector_dot.onnx) ✅<br>[`einsum_vector_dot_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_vector_dot_f64.onnx) ✅<br>[`einsum_matrix_vector`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_matrix_vector.onnx) ✅<br>[`einsum_matrix_vector_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_matrix_vector_f64.onnx) ✅<br>[`einsum_matrix_matrix_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_matrix_matrix_dynamic.onnx) ✅<br>[`einsum_matrix_matrix_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_matrix_matrix_dynamic_f64.onnx) ✅<br>[`einsum_matrix_matrix`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_matrix_matrix.onnx) ✅<br>[`einsum_matrix_matrix_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_matrix_matrix_f64.onnx) ✅<br>[`einsum_transpose`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_transpose.onnx) ✅<br>[`einsum_transpose_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_transpose_f64.onnx) ✅<br>[`einsum_batch_transpose_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_batch_transpose_dynamic.onnx) ✅<br>[`einsum_batch_transpose_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_batch_transpose_dynamic_f64.onnx) ✅<br>[`einsum_batch_transpose`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_batch_transpose.onnx) ✅<br>[`einsum_batch_transpose_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_batch_transpose_f64.onnx) ✅<br>[`einsum_diag`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_diag.onnx) ✅<br>[`einsum_diag_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_diag_f64.onnx) ✅<br>[`einsum_sum_reduce`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_sum_reduce.onnx) ✅<br>[`einsum_sum_reduce_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_sum_reduce_f64.onnx) ✅<br>[`einsum_multi_operand`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_multi_operand.onnx) ✅<br>[`einsum_multi_operand_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_multi_operand_f64.onnx) ✅<br>[`einsum_attention_logits_orig_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_orig_dynamic.onnx) ✅<br>[`einsum_attention_logits_orig_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_orig_dynamic_f64.onnx) ✅<br>[`einsum_attention_logits_orig`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_orig.onnx) ✅<br>[`einsum_attention_logits_orig_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_orig_f64.onnx) ✅<br>[`einsum_attention_output_orig_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_orig_dynamic.onnx) ✅<br>[`einsum_attention_output_orig_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_orig_dynamic_f64.onnx) ✅<br>[`einsum_attention_output_orig`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_orig.onnx) ✅<br>[`einsum_attention_output_orig_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_orig_f64.onnx) ✅<br>[`einsum_attention_logits_batched_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_batched_dynamic.onnx) ✅<br>[`einsum_attention_logits_batched_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_batched_dynamic_f64.onnx) ✅<br>[`einsum_attention_logits_batched`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_batched.onnx) ✅<br>[`einsum_attention_logits_batched_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_batched_f64.onnx) ✅<br>[`einsum_attention_output_batched_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_batched_dynamic.onnx) ✅<br>[`einsum_attention_output_batched_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_batched_dynamic_f64.onnx) ✅<br>[`einsum_attention_output_batched`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_batched.onnx) ✅<br>[`einsum_attention_output_batched_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_output_batched_f64.onnx) ✅<br>[`einsum_ellipsis_rank_mismatch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_ellipsis_rank_mismatch.onnx) ✅<br>[`einsum_ellipsis_rank_mismatch_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_ellipsis_rank_mismatch_f64.onnx) ✅<br>[`einsum_attention_logits_batched_rank_mismatch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_batched_rank_mismatch.onnx) ✅<br>[`einsum_attention_logits_batched_rank_mismatch_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/einsum_attention_logits_batched_rank_mismatch_f64.onnx) ✅ | v0.1.0 |
| [jnp.linspace](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linspace.html (Static Version for jax2onnx)) | [Constant](https://onnx.ai/onnx/operators/onnx__Constant.html) | [`linspace_static_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_basic.onnx) ✅<br>[`linspace_static_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_basic_f64.onnx) ✅<br>[`linspace_static_endpoint_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_endpoint_false.onnx) ✅<br>[`linspace_static_endpoint_false_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_endpoint_false_f64.onnx) ✅<br>[`linspace_static_num_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_num_1.onnx) ✅<br>[`linspace_static_num_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_num_1_f64.onnx) ✅<br>[`linspace_static_num_0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_num_0.onnx) ✅<br>[`linspace_static_num_0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_num_0_f64.onnx) ✅<br>[`linspace_static_int_inputs_default_dtype`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_int_inputs_default_dtype.onnx) ✅<br>[`linspace_static_int_inputs_default_dtype_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/linspace_static_int_inputs_default_dtype_f64.onnx) ✅ | v0.5.2 |
| [jnp.matmul](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.matmul.html) | [MatMul](https://onnx.ai/onnx/operators/onnx__MatMul.html) | [`matmul_2d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_2d.onnx) ✅<br>[`matmul_2d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_2d_f64.onnx) ✅<br>[`matmul_1d_2d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_1d_2d.onnx) ✅<br>[`matmul_1d_2d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_1d_2d_f64.onnx) ✅<br>[`matmul_2d_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_2d_1d.onnx) ✅<br>[`matmul_2d_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_2d_1d_f64.onnx) ✅<br>[`matmul_dynamic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_dynamic.onnx) ✅<br>[`matmul_dynamic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_dynamic_f64.onnx) ✅<br>[`matmul_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic.onnx) ✅<br>[`matmul_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_f64.onnx) ✅<br>[`matmul_dynamic_a_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_a_dynamic.onnx) ✅<br>[`matmul_dynamic_a_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_a_dynamic_f64.onnx) ✅<br>[`matmul_dynamic_a`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_a.onnx) ✅<br>[`matmul_dynamic_a_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_dynamic_a_f64.onnx) ✅<br>[`matmul_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_1d.onnx) ✅<br>[`matmul_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_1d_f64.onnx) ✅<br>[`matmul_3d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_3d.onnx) ✅<br>[`matmul_3d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/matmul_3d_f64.onnx) ✅ | v0.1.0 |
| [jnp.prod](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.prod.html#jax.numpy.prod) | [ReduceProd](https://onnx.ai/onnx/operators/onnx__ReduceProd.html) | [`basic_prod`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/basic_prod.onnx) ✅<br>[`basic_prod_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/basic_prod_f64.onnx) ✅<br>[`prod_with_axis`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/prod_with_axis.onnx) ✅<br>[`prod_with_axis_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/prod_with_axis_f64.onnx) ✅<br>[`prod_with_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/prod_with_keepdims.onnx) ✅<br>[`prod_with_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/prod_with_keepdims_f64.onnx) ✅ | v0.6.2 |
| [jnp.reshape](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.reshape.html) | [Reshape](https://onnx.ai/onnx/operators/onnx__Reshape.html) | [`reshape_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_1.onnx) ✅<br>[`reshape_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_1_f64.onnx) ✅<br>[`reshape_2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_2.onnx) ✅<br>[`reshape_2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_2_f64.onnx) ✅<br>[`reshape_3`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_3.onnx) ✅<br>[`reshape_3_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_3_f64.onnx) ✅<br>[`reshape_4_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_4_dynamic.onnx) ✅<br>[`reshape_4_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_4_dynamic_f64.onnx) ✅<br>[`reshape_4`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_4.onnx) ✅<br>[`reshape_4_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_4_f64.onnx) ✅<br>[`reshape_to_scalar`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_to_scalar.onnx) ✅<br>[`reshape_to_scalar_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_to_scalar_f64.onnx) ✅<br>[`reshape_from_scalar`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_from_scalar.onnx) ✅<br>[`reshape_from_scalar_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_from_scalar_f64.onnx) ✅<br>[`reshape_cnn_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_cnn_dynamic.onnx) ✅<br>[`reshape_cnn_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_cnn_dynamic_f64.onnx) ✅<br>[`reshape_cnn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_cnn.onnx) ✅<br>[`reshape_cnn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_cnn_f64.onnx) ✅<br>[`reshape_valid_flatten_trailing`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_valid_flatten_trailing.onnx) ✅<br>[`reshape_valid_flatten_trailing_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_valid_flatten_trailing_f64.onnx) ✅<br>[`reshape_with_target_shape_from_symbolic_dim_computation`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_with_target_shape_from_symbolic_dim_computation.onnx) ✅<br>[`reshape_with_target_shape_from_symbolic_dim_computation_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/reshape_with_target_shape_from_symbolic_dim_computation_f64.onnx) ✅ | v0.1.0 |
| [jnp.select](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.select.html) | [Where](https://onnx.ai/onnx/operators/onnx__Where.html) | [`select_simple`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_simple.onnx) ✅<br>[`select_simple_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_simple_f64.onnx) ✅<br>[`select_broadcast`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_broadcast.onnx) ✅<br>[`select_broadcast_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_broadcast_f64.onnx) ✅<br>[`select_gpt2_attention_mask_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_gpt2_attention_mask_dynamic.onnx) ✅<br>[`select_gpt2_attention_mask_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_gpt2_attention_mask_dynamic_f64.onnx) ✅<br>[`select_gpt2_attention_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_gpt2_attention_mask.onnx) ✅<br>[`select_gpt2_attention_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/select_gpt2_attention_mask_f64.onnx) ✅ | v0.7.1 |
| [jnp.shape](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.shape.html) | [Shape](https://onnx.ai/onnx/operators/onnx__Shape.html) | [`shape_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/shape_basic.onnx) ✅<br>[`shape_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/shape_basic_f64.onnx) ✅<br>[`shape_dynamic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/shape_dynamic_dynamic.onnx) ✅<br>[`shape_dynamic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/shape_dynamic_dynamic_f64.onnx) ✅<br>[`shape_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/shape_dynamic.onnx) ✅<br>[`shape_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/shape_dynamic_f64.onnx) ✅ | 0.4.0 |
| [jnp.sort](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.sort.html) | [TopK](https://onnx.ai/onnx/operators/onnx__TopK.html) | [`sort_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/sort_1d.onnx) ✅<br>[`sort_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/sort_1d_f64.onnx) ✅<br>[`sort_2d_axis0_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/sort_2d_axis0_dynamic.onnx) ✅<br>[`sort_2d_axis0_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/sort_2d_axis0_dynamic_f64.onnx) ✅<br>[`sort_2d_axis0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/sort_2d_axis0.onnx) ✅<br>[`sort_2d_axis0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/sort_2d_axis0_f64.onnx) ✅ | v0.5.2 |
| [jnp.split](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.split.html) | [Split](https://onnx.ai/onnx/operators/onnx__Split.html) | [`split_by_sections`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_sections.onnx) ✅<br>[`split_by_sections_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_sections_f64.onnx) ✅<br>[`split_by_indices`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_indices.onnx) ✅<br>[`split_by_indices_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_indices_f64.onnx) ✅<br>[`split_by_indices_symbolic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_indices_symbolic_dynamic.onnx) ✅<br>[`split_by_indices_symbolic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_indices_symbolic_dynamic_f64.onnx) ✅<br>[`split_by_indices_symbolic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_indices_symbolic.onnx) ✅<br>[`split_by_indices_symbolic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/split_by_indices_symbolic_f64.onnx) ✅ | v0.7.2 |
| [jnp.squeeze](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.squeeze.html) | [Squeeze](https://onnx.ai/onnx/operators/onnx__Squeeze.html) | [`squeeze_single_dim`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_single_dim.onnx) ✅<br>[`squeeze_single_dim_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_single_dim_f64.onnx) ✅<br>[`squeeze_multiple_dims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_multiple_dims.onnx) ✅<br>[`squeeze_multiple_dims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_multiple_dims_f64.onnx) ✅<br>[`squeeze_vit_output`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_vit_output.onnx) ✅<br>[`squeeze_vit_output_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_vit_output_f64.onnx) ✅<br>[`squeeze_dynamic_batch_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_batch_dynamic.onnx) ✅<br>[`squeeze_dynamic_batch_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_batch_dynamic_f64.onnx) ✅<br>[`squeeze_dynamic_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_batch.onnx) ✅<br>[`squeeze_dynamic_batch_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_batch_f64.onnx) ✅<br>[`squeeze_all_dims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_all_dims.onnx) ✅<br>[`squeeze_all_dims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_all_dims_f64.onnx) ✅<br>[`squeeze_negative_axis`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_negative_axis.onnx) ✅<br>[`squeeze_negative_axis_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_negative_axis_f64.onnx) ✅<br>[`squeeze_negative_axis_tuple`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_negative_axis_tuple.onnx) ✅<br>[`squeeze_negative_axis_tuple_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_negative_axis_tuple_f64.onnx) ✅<br>[`squeeze_dynamic_and_negative_axis_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_and_negative_axis_dynamic.onnx) ✅<br>[`squeeze_dynamic_and_negative_axis_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_and_negative_axis_dynamic_f64.onnx) ✅<br>[`squeeze_dynamic_and_negative_axis`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_and_negative_axis.onnx) ✅<br>[`squeeze_dynamic_and_negative_axis_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/squeeze_dynamic_and_negative_axis_f64.onnx) ✅ | v0.1.0 |
| [jnp.stack](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.stack.html) | [Concat](https://onnx.ai/onnx/operators/onnx__Concat.html)<br>[Unsqueeze](https://onnx.ai/onnx/operators/onnx__Unsqueeze.html) | [`stack_axis_0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_axis_0.onnx) ✅<br>[`stack_axis_0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_axis_0_f64.onnx) ✅<br>[`stack_axis_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_axis_1.onnx) ✅<br>[`stack_axis_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_axis_1_f64.onnx) ✅<br>[`stack_negative_axis`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_negative_axis.onnx) ✅<br>[`stack_negative_axis_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_negative_axis_f64.onnx) ✅<br>[`stack_scalars`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_scalars.onnx) ✅<br>[`stack_scalars_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/stack_scalars_f64.onnx) ✅ | v0.7.1 |
| [jnp.take](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.take.html) | [Gather](https://onnx.ai/onnx/operators/onnx__Gather.html) | [`take_data_dependent_indices`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/take_data_dependent_indices.onnx) ✅ | v0.7.0 |
| [jnp.tile](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.tile.html) | [Tile](https://onnx.ai/onnx/operators/onnx__Tile.html) | [`tile_repeats`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_repeats.onnx) ✅<br>[`tile_repeats_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_repeats_f64.onnx) ✅<br>[`tile_a`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_a.onnx) ✅<br>[`tile_a_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_a_f64.onnx) ✅<br>[`tile_b`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_b.onnx) ✅<br>[`tile_b_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_b_f64.onnx) ✅<br>[`tile_c`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_c.onnx) ✅<br>[`tile_c_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_c_f64.onnx) ✅<br>[`tile_d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_d.onnx) ✅<br>[`tile_d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_d_f64.onnx) ✅<br>[`tile_dynamic_input_static`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_dynamic_input_static.onnx) ✅<br>[`tile_dynamic_input_static_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_dynamic_input_static_f64.onnx) ✅<br>[`tile_dynamic_input_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_dynamic_input_dynamic.onnx) ✅<br>[`tile_dynamic_input_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_dynamic_input_dynamic_f64.onnx) ✅<br>[`tile_dynamic_input`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_dynamic_input.onnx) ✅<br>[`tile_dynamic_input_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_dynamic_input_f64.onnx) ✅<br>[`tile_pad`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_pad.onnx) ✅<br>[`tile_pad_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_pad_f64.onnx) ✅<br>[`tile_with_symbolic_repeats_static`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_with_symbolic_repeats_static.onnx) ✅<br>[`tile_with_symbolic_repeats_static_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_with_symbolic_repeats_static_f64.onnx) ✅<br>[`tile_with_symbolic_repeats_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_with_symbolic_repeats_dynamic.onnx) ✅<br>[`tile_with_symbolic_repeats_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_with_symbolic_repeats_dynamic_f64.onnx) ✅<br>[`tile_with_symbolic_repeats`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_with_symbolic_repeats.onnx) ✅<br>[`tile_with_symbolic_repeats_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_with_symbolic_repeats_f64.onnx) ✅<br>[`tile_param_symbolic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_param_symbolic_dynamic.onnx) ✅<br>[`tile_param_symbolic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_param_symbolic_dynamic_f64.onnx) ✅<br>[`tile_param_symbolic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_param_symbolic.onnx) ✅<br>[`tile_param_symbolic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/tile_param_symbolic_f64.onnx) ✅ | v0.1.0 |
| [jnp.transpose](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.transpose.html) | [Transpose](https://onnx.ai/onnx/operators/onnx__Transpose.html) | [`transpose_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_basic.onnx) ✅<br>[`transpose_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_basic_f64.onnx) ✅<br>[`transpose_reverse`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_reverse.onnx) ✅<br>[`transpose_reverse_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_reverse_f64.onnx) ✅<br>[`transpose_4d_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_4d_dynamic.onnx) ✅<br>[`transpose_4d_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_4d_dynamic_f64.onnx) ✅<br>[`transpose_4d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_4d.onnx) ✅<br>[`transpose_4d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_4d_f64.onnx) ✅<br>[`transpose_square_matrix`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_square_matrix.onnx) ✅<br>[`transpose_square_matrix_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_square_matrix_f64.onnx) ✅<br>[`transpose_high_dim`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_high_dim.onnx) ✅<br>[`transpose_high_dim_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_high_dim_f64.onnx) ✅<br>[`transpose_no_axes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_no_axes.onnx) ✅<br>[`transpose_no_axes_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_no_axes_f64.onnx) ✅<br>[`transpose_3d_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_3d_dynamic.onnx) ✅<br>[`transpose_3d_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_3d_dynamic_f64.onnx) ✅<br>[`transpose_3d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_3d.onnx) ✅<br>[`transpose_3d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/transpose_3d_f64.onnx) ✅ | v0.1.0 |
| [jnp.unstack](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.unstack.html) | [Split](https://onnx.ai/onnx/operators/onnx__Split.html)<br>[Squeeze](https://onnx.ai/onnx/operators/onnx__Squeeze.html) | [`unstack_axis_0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/unstack_axis_0.onnx) ✅<br>[`unstack_axis_0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/unstack_axis_0_f64.onnx) ✅<br>[`unstack_axis_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/unstack_axis_1.onnx) ✅<br>[`unstack_axis_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/unstack_axis_1_f64.onnx) ✅<br>[`unstack_negative_axis`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/unstack_negative_axis.onnx) ✅<br>[`unstack_negative_axis_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/unstack_negative_axis_f64.onnx) ✅ | v0.7.1 |
| [jnp.where](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.where.html) | [Where](https://onnx.ai/onnx/operators/onnx__Where.html) | [`where_gpt_mask_scores_literal_else_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_literal_else_dynamic.onnx) ✅<br>[`where_gpt_mask_scores_literal_else_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_literal_else_dynamic_f64.onnx) ✅<br>[`where_gpt_mask_scores_literal_else`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_literal_else.onnx) ✅<br>[`where_gpt_mask_scores_literal_else_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_literal_else_f64.onnx) ✅<br>[`where_simple`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_simple.onnx) ✅<br>[`where_simple_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_simple_f64.onnx) ✅<br>[`where_broadcast`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_broadcast.onnx) ✅<br>[`where_broadcast_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_broadcast_f64.onnx) ✅<br>[`where_multidim_condition_scalar_branches_broadcast`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_multidim_condition_scalar_branches_broadcast.onnx) ✅<br>[`where_multidim_condition_scalar_branches_broadcast_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_multidim_condition_scalar_branches_broadcast_f64.onnx) ✅<br>[`where_multidim_condition_scalar_branches_broadcast`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_multidim_condition_scalar_branches_broadcast.onnx) ✅<br>[`where_multidim_condition_scalar_branches_broadcast_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_multidim_condition_scalar_branches_broadcast_f64.onnx) ✅<br>[`where_A`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_A.onnx) ✅<br>[`where_A_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_A_f64.onnx) ✅<br>[`where_B`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_B.onnx) ✅<br>[`where_B_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_B_f64.onnx) ✅<br>[`where_gpt_mask_scores_scalar_else_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_scalar_else_dynamic.onnx) ✅<br>[`where_gpt_mask_scores_scalar_else_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_scalar_else_dynamic_f64.onnx) ✅<br>[`where_gpt_mask_scores_scalar_else`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_scalar_else.onnx) ✅<br>[`where_gpt_mask_scores_scalar_else_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_gpt_mask_scores_scalar_else_f64.onnx) ✅<br>[`where_int_condition_cast`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_int_condition_cast.onnx) ✅<br>[`where_int_condition_cast_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_int_condition_cast_f64.onnx) ✅<br>[`where_literal_else_pyfloat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_literal_else_pyfloat.onnx) ✅<br>[`where_literal_else_pyfloat_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_literal_else_pyfloat_f64.onnx) ✅<br>[`where_jax_int_literals_broadcast_f64_mode`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_jax_int_literals_broadcast_f64_mode.onnx) ✅<br>[`where_simple`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_simple.onnx) ✅<br>[`where_simple_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/jnp/where_simple_f64.onnx) ✅ | v0.5.2 |
| [lax.abs](https://docs.jax.dev/en/latest/_autosummary/jax.lax.abs.html) | [Abs](https://onnx.ai/onnx/operators/onnx__Abs.html) | [`abs`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/abs.onnx) ✅<br>[`abs_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/abs_f64.onnx) ✅ | v0.5.0 |
| [lax.add](https://docs.jax.dev/en/latest/_autosummary/jax.lax.add.html) | [Add](https://onnx.ai/onnx/operators/onnx__Add.html) | [`add`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/add.onnx) ✅<br>[`add_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/add_f64.onnx) ✅ | v0.2.0 |
| [lax.and](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.html#jax.lax.bitwise_and) | [And](https://onnx.ai/onnx/operators/onnx__And.html)<br>[BitwiseAnd](https://onnx.ai/onnx/operators/onnx__BitwiseAnd.html) | [`and_bool`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/and_bool.onnx) ✅<br>[`and_bool_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/and_bool_f64.onnx) ✅<br>[`and_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/and_int.onnx) ✅<br>[`and_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/and_int_f64.onnx) ✅ | v0.6.5 |
| [lax.argmax](https://docs.jax.dev/en/latest/_autosummary/jax.lax.argmax.html) | [ArgMax](https://onnx.ai/onnx/operators/onnx__ArgMax.html) | [`argmax_float_axis0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_float_axis0.onnx) ✅<br>[`argmax_float_axis0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_float_axis0_f64.onnx) ✅<br>[`argmax_float_axis1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_float_axis1.onnx) ✅<br>[`argmax_float_axis1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_float_axis1_f64.onnx) ✅<br>[`argmax_boolean_input_axis0_specific_values`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_boolean_input_axis0_specific_values.onnx) ✅<br>[`argmax_boolean_input_axis0_specific_values_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_boolean_input_axis0_specific_values_f64.onnx) ✅<br>[`argmax_boolean_input_axis1_specific_values`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_boolean_input_axis1_specific_values.onnx) ✅<br>[`argmax_boolean_input_axis1_specific_values_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_boolean_input_axis1_specific_values_f64.onnx) ✅<br>[`argmax_boolean_random_input_axis0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_boolean_random_input_axis0.onnx) ✅<br>[`argmax_boolean_random_input_axis0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmax_boolean_random_input_axis0_f64.onnx) ✅ | v0.2.0 |
| [lax.argmin](https://docs.jax.dev/en/latest/_autosummary/jax.lax.argmin.html) | [ArgMin](https://onnx.ai/onnx/operators/onnx__ArgMin.html) | [`argmin_test1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmin_test1.onnx) ✅<br>[`argmin_test1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmin_test1_f64.onnx) ✅<br>[`argmin_test2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmin_test2.onnx) ✅<br>[`argmin_test2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/argmin_test2_f64.onnx) ✅ | v0.2.0 |
| [lax.broadcast_in_dim](https://jax.readthedocs.io/en/latest/jax-primitives.html) | [Expand](https://onnx.ai/onnx/operators/onnx__Expand.html)<br>[Identity](https://onnx.ai/onnx/operators/onnx__Identity.html)<br>[Reshape](https://onnx.ai/onnx/operators/onnx__Reshape.html) | [`broadcast_in_dim`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim.onnx) ✅<br>[`broadcast_in_dim_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_f64.onnx) ✅<br>[`broadcast_in_dim_2d_to_3d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_2d_to_3d.onnx) ✅<br>[`broadcast_in_dim_2d_to_3d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_2d_to_3d_f64.onnx) ✅<br>[`broadcast_in_dim_scalar`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_scalar.onnx) ✅<br>[`broadcast_in_dim_scalar_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_scalar_f64.onnx) ✅<br>[`broadcast_in_dim_batch_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_batch_dynamic.onnx) ✅<br>[`broadcast_in_dim_batch_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_batch_dynamic_f64.onnx) ✅<br>[`broadcast_in_dim_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_batch.onnx) ✅<br>[`broadcast_in_dim_batch_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_batch_f64.onnx) ✅<br>[`broadcast_in_dim_dynamic_B_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_dynamic_B_dynamic.onnx) ✅<br>[`broadcast_in_dim_dynamic_B_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_dynamic_B_dynamic_f64.onnx) ✅<br>[`broadcast_in_dim_dynamic_B`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_dynamic_B.onnx) ✅<br>[`broadcast_in_dim_dynamic_B_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcast_in_dim_dynamic_B_f64.onnx) ✅ | v0.2.0 |
| [lax.concatenate](https://docs.jax.dev/en/latest/_autosummary/jax.lax.concatenate.html) | [Concat](https://onnx.ai/onnx/operators/onnx__Concat.html) | [`concatenate`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate.onnx) ✅<br>[`concatenate_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_f64.onnx) ✅<br>[`concatenate_axis1_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_axis1_dynamic.onnx) ✅<br>[`concatenate_axis1_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_axis1_dynamic_f64.onnx) ✅<br>[`concatenate_axis1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_axis1.onnx) ✅<br>[`concatenate_axis1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_axis1_f64.onnx) ✅<br>[`concatenate_axis0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_axis0.onnx) ✅<br>[`concatenate_axis0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_axis0_f64.onnx) ✅<br>[`concatenate_3d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_3d.onnx) ✅<br>[`concatenate_3d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/concatenate_3d_f64.onnx) ✅ | v0.2.0 |
| [lax.cond](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.cond.html) | [If](https://onnx.ai/onnx/operators/onnx__If.html) | [`cond_scalar`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_scalar.onnx) ✅<br>[`cond_scalar_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_scalar_f64.onnx) ✅<br>[`cond_multiple_operands_in_tuple`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_multiple_operands_in_tuple.onnx) ✅<br>[`cond_multiple_operands_in_tuple_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_multiple_operands_in_tuple_f64.onnx) ✅<br>[`cond_my_new_complex_scenario`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_my_new_complex_scenario.onnx) ✅<br>[`cond_my_new_complex_scenario_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_my_new_complex_scenario_f64.onnx) ✅<br>[`cond_nested_conditional`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_nested_conditional.onnx) ✅<br>[`cond_nested_conditional_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_nested_conditional_f64.onnx) ✅<br>[`cond_variables`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_variables.onnx) ✅<br>[`cond_variables_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_variables_f64.onnx) ✅<br>[`cond_internal_constant_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_internal_constant_f64.onnx) ✅<br>[`cond_passthrough_identity`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_passthrough_identity.onnx) ✅<br>[`cond_passthrough_identity_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_passthrough_identity_f64.onnx) ✅<br>[`cond_with_scatter`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_with_scatter.onnx) ✅<br>[`cond_with_scatter_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cond_with_scatter_f64.onnx) ✅ | v0.5.1 |
| [lax.conv](https://docs.jax.dev/en/latest/_autosummary/jax.lax.conv.html) | [Conv](https://onnx.ai/onnx/operators/onnx__Conv.html) | [`conv`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/conv.onnx) ✅<br>[`conv2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/conv2.onnx) ✅<br>[`conv_general_dilated_nhwc_output`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/conv_general_dilated_nhwc_output.onnx) ✅ | v0.2.0 |
| [lax.convert_element_type](https://docs.jax.dev/en/latest/_autosummary/jax.lax.convert_element_type.html) | [Cast](https://onnx.ai/onnx/operators/onnx__Cast.html) | [`convert_element_type`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/convert_element_type.onnx) ✅<br>[`convert_element_type_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/convert_element_type_f64.onnx) ✅ | v0.2.0 |
| [lax.copy](Handles the JAX primitive lax.copy_p. Note: jax.lax.copy API is removed.) | [Identity](https://onnx.ai/onnx/operators/onnx__Identity.html) | [`copy_float32_array`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/copy_float32_array.onnx) ✅<br>[`copy_int64_scalar`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/copy_int64_scalar.onnx) ✅ | <your_current_version> |
| [lax.cos](https://docs.jax.dev/en/latest/_autosummary/jax.lax.cos.html) | [Cos](https://onnx.ai/onnx/operators/onnx__Cos.html) | [`cos`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cos.onnx) ✅<br>[`cos_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cos_f64.onnx) ✅ | v0.4.4 |
| [lax.cosh](https://docs.jax.dev/en/latest/_autosummary/jax.lax.cosh.html) | [Cosh](https://onnx.ai/onnx/operators/onnx__Cosh.html) | [`cosh`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cosh.onnx) ✅<br>[`cosh_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/cosh_f64.onnx) ✅ | v0.4.4 |
| [lax.device_put](https://docs.jax.dev/en/latest/_autosummary/jax.device_put.html#jax.device_put) | [Identity](https://onnx.ai/onnx/operators/onnx__Identity.html) | [`device_put_array`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/device_put_array.onnx) ✅<br>[`device_put_array_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/device_put_array_f64.onnx) ✅<br>[`device_put_scalar`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/device_put_scalar.onnx) ✅<br>[`device_put_scalar_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/device_put_scalar_f64.onnx) ✅ | v0.4.0 |
| [lax.div](https://docs.jax.dev/en/latest/_autosummary/jax.lax.div.html) | [Div](https://onnx.ai/onnx/operators/onnx__Div.html) | [`div`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/div.onnx) ✅<br>[`div_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/div_f64.onnx) ✅ | v0.2.0 |
| [lax.dot_general](https://docs.jax.dev/en/latest/_autosummary/jax.lax.dot_general.html) | [MatMul](https://onnx.ai/onnx/operators/onnx__MatMul.html) | [`dot_general`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dot_general.onnx) ✅<br>[`dot_general_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dot_general_f64.onnx) ✅<br>[`dot_general_lhs1_rhs1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dot_general_lhs1_rhs1.onnx) ✅<br>[`dot_general_lhs1_rhs1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dot_general_lhs1_rhs1_f64.onnx) ✅ | v0.2.0 |
| [lax.dynamic_slice](https://docs.jax.dev/en/latest/_autosummary/jax.lax.dynamic_slice.html) | [Slice](https://onnx.ai/onnx/operators/onnx__Slice.html) | [`dynamic_slice_test1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_test1.onnx) ✅<br>[`dynamic_slice_test1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_test1_f64.onnx) ✅<br>[`dynamic_slice_2d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_2d.onnx) ✅<br>[`dynamic_slice_2d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_2d_f64.onnx) ✅<br>[`dynamic_slice_3d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_3d.onnx) ✅<br>[`dynamic_slice_3d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_3d_f64.onnx) ✅<br>[`dynamic_slice_vit_like_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_vit_like_dynamic.onnx) ✅<br>[`dynamic_slice_vit_like_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_vit_like_dynamic_f64.onnx) ✅<br>[`dynamic_slice_vit_like`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_vit_like.onnx) ✅<br>[`dynamic_slice_vit_like_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/dynamic_slice_vit_like_f64.onnx) ✅ | v0.1.0 |
| [lax.eq](https://docs.jax.dev/en/latest/_autosummary/jax.lax.eq.html) | [Equal](https://onnx.ai/onnx/operators/onnx__Equal.html) | [`eq`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/eq.onnx) ✅<br>[`eq_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/eq_f64.onnx) ✅ | v0.2.0 |
| [lax.exp](https://docs.jax.dev/en/latest/_autosummary/jax.lax.exp.html) | [Exp](https://onnx.ai/onnx/operators/onnx__Exp.html) | [`exp`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/exp.onnx) ✅<br>[`exp_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/exp_f64.onnx) ✅ | v0.2.0 |
| [lax.fori_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html) | [Loop](https://onnx.ai/onnx/operators/onnx__Loop.html) | [`fori_loop_counter`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_counter.onnx) ✅<br>[`fori_loop_counter_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_counter_f64.onnx) ✅<br>[`fori_loop_zero`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_zero.onnx) ✅<br>[`fori_loop_zero_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_zero_f64.onnx) ✅<br>[`fori_loop_vector`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_vector.onnx) ✅<br>[`fori_loop_vector_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_vector_f64.onnx) ✅<br>[`fori_loop_example`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_example.onnx) ✅<br>[`fori_loop_example_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_example_f64.onnx) ✅<br>[`fori_loop_test`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_test.onnx) ✅<br>[`fori_loop_test_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/fori_loop_test_f64.onnx) ➖ | v0.5.1 |
| [lax.gather](https://docs.jax.dev/en/latest/_autosummary/jax.lax.gather.html) | [GatherND](https://onnx.ai/onnx/operators/onnx__GatherND.html) | [`gather_static`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gather_static.onnx) ✅<br>[`gather_static_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gather_static_f64.onnx) ✅<br>[`gather_dynamic_batch_simple_index_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gather_dynamic_batch_simple_index_dynamic.onnx) ✅<br>[`gather_dynamic_batch_simple_index_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gather_dynamic_batch_simple_index_dynamic_f64.onnx) ✅<br>[`gather_dynamic_batch_simple_index`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gather_dynamic_batch_simple_index.onnx) ✅<br>[`gather_dynamic_batch_simple_index_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gather_dynamic_batch_simple_index_f64.onnx) ✅ | v0.2.0 |
| [lax.gt](https://docs.jax.dev/en/latest/_autosummary/jax.lax.gt.html) | [Greater](https://onnx.ai/onnx/operators/onnx__Greater.html) | [`gt`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gt.onnx) ✅<br>[`gt_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/gt_f64.onnx) ✅ | v0.2.0 |
| [lax.integer_pow](https://docs.jax.dev/en/latest/_autosummary/jax.lax.integer_pow.html) | [Pow](https://onnx.ai/onnx/operators/onnx__Pow.html) | [`integer_pow`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/integer_pow.onnx) ✅<br>[`integer_pow_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/integer_pow_f64.onnx) ✅ | v0.2.0 |
| [lax.iota](https://docs.jax.dev/en/latest/_autosummary/jax.lax.iota.html) | [Range](https://onnx.ai/onnx/operators/onnx__Range.html) | [`iota_int32`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/iota_int32.onnx) ✅<br>[`iota_int32_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/iota_int32_f64.onnx) ✅<br>[`iota_float32`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/iota_float32.onnx) ✅<br>[`iota_float32_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/iota_float32_f64.onnx) ✅<br>[`broadcasted_iota`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcasted_iota.onnx) ✅<br>[`broadcasted_iota_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/broadcasted_iota_f64.onnx) ✅ | v0.5.0 |
| [lax.log](https://docs.jax.dev/en/latest/_autosummary/jax.lax.log.html) | [Log](https://onnx.ai/onnx/operators/onnx__Log.html) | [`log`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/log.onnx) ✅<br>[`log_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/log_f64.onnx) ✅ | v0.2.0 |
| [lax.logistic](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.logistic.html) | [Sigmoid](https://onnx.ai/onnx/operators/onnx__Sigmoid.html) | [`lax_logistic_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_logistic_basic.onnx) ✅<br>[`lax_logistic_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_logistic_basic_f64.onnx) ✅ | v0.7.2 |
| [lax.lt](https://docs.jax.dev/en/latest/_autosummary/jax.lax.lt.html) | [Less](https://onnx.ai/onnx/operators/onnx__Less.html) | [`lt`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lt.onnx) ✅<br>[`lt_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lt_f64.onnx) ✅ | v0.2.0 |
| [lax.max](https://docs.jax.dev/en/latest/_autosummary/jax.lax.max.html) | [Max](https://onnx.ai/onnx/operators/onnx__Max.html) | [`max`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/max.onnx) ✅<br>[`max_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/max_f64.onnx) ✅ | v0.2.0 |
| [lax.min](https://docs.jax.dev/en/latest/_autosummary/jax.lax.min.html) | [Min](https://onnx.ai/onnx/operators/onnx__Min.html) | [`min_test1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/min_test1.onnx) ✅<br>[`min_test1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/min_test1_f64.onnx) ✅ | v0.1.0 |
| [lax.mul](https://docs.jax.dev/en/latest/_autosummary/jax.lax.mul.html) | [Mul](https://onnx.ai/onnx/operators/onnx__Mul.html) | [`mul_test1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/mul_test1.onnx) ✅<br>[`mul_test1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/mul_test1_f64.onnx) ✅<br>[`mul_test2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/mul_test2.onnx) ✅<br>[`mul_test2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/mul_test2_f64.onnx) ✅ | v0.1.0 |
| [lax.ne](https://docs.jax.dev/en/latest/_autosummary/jax.lax.ne.html) | [Equal](https://onnx.ai/onnx/operators/onnx__Equal.html)<br>[Not](https://onnx.ai/onnx/operators/onnx__Not.html) | [`ne`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/ne.onnx) ✅<br>[`ne_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/ne_f64.onnx) ✅ | v0.2.0 |
| [lax.neg](https://docs.jax.dev/en/latest/_autosummary/jax.lax.neg.html) | [Neg](https://onnx.ai/onnx/operators/onnx__Neg.html) | [`neg`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/neg.onnx) ✅<br>[`neg_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/neg_f64.onnx) ✅ | v0.2.0 |
| [lax.reduce_and](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.all.html) | [Cast](https://onnx.ai/onnx/operators/onnx__Cast.html)<br>[ReduceMin](https://onnx.ai/onnx/operators/onnx__ReduceMin.html) | [`reduce_and_all_true`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_and_all_true.onnx) ✅<br>[`reduce_and_all_true_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_and_all_true_f64.onnx) ✅<br>[`reduce_and_one_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_and_one_false.onnx) ✅<br>[`reduce_and_one_false_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_and_one_false_f64.onnx) ✅<br>[`reduce_and_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_and_keepdims.onnx) ✅<br>[`reduce_and_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_and_keepdims_f64.onnx) ✅ | v0.6.1 |
| [lax.reduce_max](https://docs.jax.dev/en/latest/_autosummary/jax.lax.reduce_max.html) | [ReduceMax](https://onnx.ai/onnx/operators/onnx__ReduceMax.html) | [`reduce_max`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_max.onnx) ✅<br>[`reduce_max_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_max_f64.onnx) ✅<br>[`reduce_max_allaxes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_max_allaxes.onnx) ✅<br>[`reduce_max_allaxes_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_max_allaxes_f64.onnx) ✅<br>[`reduce_max_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_max_keepdims.onnx) ✅<br>[`reduce_max_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_max_keepdims_f64.onnx) ✅ | v0.2.0 |
| [lax.reduce_min](https://docs.jax.dev/en/latest/_autosummary/jax.lax.reduce_min.html) | [ReduceMin](https://onnx.ai/onnx/operators/onnx__ReduceMin.html) | [`reduce_min`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_min.onnx) ✅<br>[`reduce_min_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_min_f64.onnx) ✅<br>[`reduce_min_allaxes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_min_allaxes.onnx) ✅<br>[`reduce_min_allaxes_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_min_allaxes_f64.onnx) ✅<br>[`reduce_min_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_min_keepdims.onnx) ✅<br>[`reduce_min_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_min_keepdims_f64.onnx) ✅ | v0.2.0 |
| [lax.reduce_or](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.any.html) | [Cast](https://onnx.ai/onnx/operators/onnx__Cast.html)<br>[ReduceMax](https://onnx.ai/onnx/operators/onnx__ReduceMax.html) | [`reduce_or_all_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_or_all_false.onnx) ✅<br>[`reduce_or_all_false_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_or_all_false_f64.onnx) ✅<br>[`reduce_or_one_true`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_or_one_true.onnx) ✅<br>[`reduce_or_one_true_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_or_one_true_f64.onnx) ✅<br>[`reduce_or_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_or_keepdims.onnx) ✅<br>[`reduce_or_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_or_keepdims_f64.onnx) ✅ | v0.6.1 |
| [lax.reduce_prod](https://docs.jax.dev/en/latest/_autosummary/jax.lax.reduce_prod.html) | [ReduceProd](https://onnx.ai/onnx/operators/onnx__ReduceProd.html) | [`reduce_prod`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod.onnx) ✅<br>[`reduce_prod_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_f64.onnx) ✅<br>[`reduce_prod_allaxes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_allaxes.onnx) ✅<br>[`reduce_prod_allaxes_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_allaxes_f64.onnx) ✅<br>[`reduce_prod_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_keepdims.onnx) ✅<br>[`reduce_prod_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_keepdims_f64.onnx) ✅<br>[`reduce_prod_dtype_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_dtype_f64.onnx) ✅<br>[`reduce_prod_dtype`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_prod_dtype.onnx) ✅ | v0.6.1 |
| [lax.reduce_sum](https://docs.jax.dev/en/latest/_autosummary/jax.lax.reduce_sum.html) | [ReduceSum](https://onnx.ai/onnx/operators/onnx__ReduceSum.html) | [`reduce_sum`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum.onnx) ✅<br>[`reduce_sum_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_f64.onnx) ✅<br>[`reduce_sum_allaxes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_allaxes.onnx) ✅<br>[`reduce_sum_allaxes_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_allaxes_f64.onnx) ✅<br>[`reduce_sum_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_keepdims.onnx) ✅<br>[`reduce_sum_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_keepdims_f64.onnx) ✅<br>[`reduce_sum_dtype_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_dtype_f64.onnx) ✅<br>[`reduce_sum_dtype`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_sum_dtype.onnx) ✅ | v0.2.0 |
| [lax.reduce_xor](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.logical_xor.html) | [Cast](https://onnx.ai/onnx/operators/onnx__Cast.html)<br>[Mod](https://onnx.ai/onnx/operators/onnx__Mod.html)<br>[ReduceSum](https://onnx.ai/onnx/operators/onnx__ReduceSum.html) | [`reduce_xor_all_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_all_false.onnx) ✅<br>[`reduce_xor_all_false_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_all_false_f64.onnx) ✅<br>[`reduce_xor_one_true`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_one_true.onnx) ✅<br>[`reduce_xor_one_true_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_one_true_f64.onnx) ✅<br>[`reduce_xor_two_true`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_two_true.onnx) ✅<br>[`reduce_xor_two_true_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_two_true_f64.onnx) ✅<br>[`reduce_xor_keepdims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_keepdims.onnx) ✅<br>[`reduce_xor_keepdims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reduce_xor_keepdims_f64.onnx) ✅ | v0.6.1 |
| [lax.rem](https://docs.jax.dev/en/latest/_autosummary/jax.lax.rem.html) | [Div](https://onnx.ai/onnx/operators/onnx__Div.html)<br>[Mod](https://onnx.ai/onnx/operators/onnx__Mod.html) | [`rem_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_int.onnx) ✅<br>[`rem_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_int_f64.onnx) ✅<br>[`rem_float`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_float.onnx) ✅<br>[`rem_float_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_float_f64.onnx) ✅<br>[`rem_int_neg`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_int_neg.onnx) ✅<br>[`rem_int_neg_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_int_neg_f64.onnx) ✅<br>[`rem_float_neg`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_float_neg.onnx) ✅<br>[`rem_float_neg_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/rem_float_neg_f64.onnx) ✅ | v0.6.5 |
| [lax.reshape](https://docs.jax.dev/en/latest/_autosummary/jax.lax.reshape.html) | [Reshape](https://onnx.ai/onnx/operators/onnx__Reshape.html) | [`reshape`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape.onnx) ✅<br>[`reshape_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_f64.onnx) ✅<br>[`reshape_valid_squeeze_middle_dim_from_problematic_source`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_valid_squeeze_middle_dim_from_problematic_source.onnx) ✅<br>[`reshape_valid_squeeze_middle_dim_from_problematic_source_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_valid_squeeze_middle_dim_from_problematic_source_f64.onnx) ✅<br>[`reshape_valid_flatten_trailing`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_valid_flatten_trailing.onnx) ✅<br>[`reshape_valid_flatten_trailing_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_valid_flatten_trailing_f64.onnx) ✅<br>[`reshape_with_target_shape_from_symbolic_dim_computation`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_target_shape_from_symbolic_dim_computation.onnx) ✅<br>[`reshape_with_target_shape_from_symbolic_dim_computation_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_target_shape_from_symbolic_dim_computation_f64.onnx) ✅<br>[`reshape_with_inferred_dimension_from_input_dynamic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_inferred_dimension_from_input_dynamic_dynamic.onnx) ✅<br>[`reshape_with_inferred_dimension_from_input_dynamic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_inferred_dimension_from_input_dynamic_dynamic_f64.onnx) ✅<br>[`reshape_with_inferred_dimension_from_input_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_inferred_dimension_from_input_dynamic.onnx) ✅<br>[`reshape_with_inferred_dimension_from_input_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_inferred_dimension_from_input_dynamic_f64.onnx) ✅<br>[`reshape_with_inferred_dimension_from_input`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_inferred_dimension_from_input.onnx) ✅<br>[`reshape_with_inferred_dimension_from_input_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/reshape_with_inferred_dimension_from_input_f64.onnx) ✅ | v0.2.0 |
| [lax.scan](https://docs.jax.dev/en/latest/_autosummary/jax.lax.scan.html) | [Scan](https://onnx.ai/onnx/operators/onnx__Scan.html) | [`scan_cumsum`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_cumsum.onnx) ✅<br>[`scan_cumsum_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_cumsum_f64.onnx) ✅<br>[`scan_carry_only`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_carry_only.onnx) ✅<br>[`scan_carry_only_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_carry_only_f64.onnx) ✅<br>[`scan_multiple_sequences`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_multiple_sequences.onnx) ✅<br>[`scan_multiple_sequences_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_multiple_sequences_f64.onnx) ✅<br>[`scan_multiple_carry`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_multiple_carry.onnx) ✅<br>[`scan_multiple_carry_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_multiple_carry_f64.onnx) ✅<br>[`scan_matrix_carry_multidim_xs`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_matrix_carry_multidim_xs.onnx) ✅<br>[`scan_matrix_carry_multidim_xs_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_matrix_carry_multidim_xs_f64.onnx) ✅<br>[`scan_no_xs`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_no_xs.onnx) ✅<br>[`scan_no_xs_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_no_xs_f64.onnx) ✅<br>[`scan_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_fn.onnx) ✅<br>[`scan_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_fn_f64.onnx) ✅<br>[`scan_jit_no_xs`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_jit_no_xs.onnx) ✅<br>[`scan_jit_no_xs_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scan_jit_no_xs_f64.onnx) ✅ | v0.5.1 |
| [lax.scatter](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scatter.html) | [ScatterND](https://onnx.ai/onnx/operators/onnx__ScatterND.html) | [`scatter_set_axis0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_set_axis0.onnx) ✅<br>[`scatter_set_axis0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_set_axis0_f64.onnx) ✅<br>[`scatter_set_middle`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_set_middle.onnx) ✅<br>[`scatter_set_middle_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_set_middle_f64.onnx) ✅<br>[`scatter_correct_axis_determination`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_correct_axis_determination.onnx) ✅<br>[`scatter_correct_axis_determination_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_correct_axis_determination_f64.onnx) ✅<br>[`scatter_updates_slice_needed_axis0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_updates_slice_needed_axis0.onnx) ✅<br>[`scatter_updates_slice_needed_axis0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_updates_slice_needed_axis0_f64.onnx) ✅<br>[`scatter_from_user_warning_shapes_valid_jax`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_from_user_warning_shapes_valid_jax.onnx) ✅<br>[`scatter_from_user_warning_shapes_valid_jax_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_from_user_warning_shapes_valid_jax_f64.onnx) ✅<br>[`scatter_user_error_scenario_precise`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_user_error_scenario_precise.onnx) ✅<br>[`scatter_user_error_scenario_precise_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_user_error_scenario_precise_f64.onnx) ✅ | v0.4.4 |
| [lax.scatter_add](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scatter_add.html) | [ScatterND](https://onnx.ai/onnx/operators/onnx__ScatterND.html) | [`scatter_add_simple_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_simple_1d.onnx) ✅<br>[`scatter_add_simple_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_simple_1d_f64.onnx) ✅<br>[`scatter_add_window_2d_operand_1d_indices`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_window_2d_operand_1d_indices.onnx) ✅<br>[`scatter_add_window_2d_operand_1d_indices_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_window_2d_operand_1d_indices_f64.onnx) ✅<br>[`scatter_add_batch_updates_1d_operand`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_batch_updates_1d_operand.onnx) ✅<br>[`scatter_add_batch_updates_1d_operand_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_batch_updates_1d_operand_f64.onnx) ✅<br>[`scatter_add_mismatched_window_dims_from_user_report`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_mismatched_window_dims_from_user_report.onnx) ✅<br>[`scatter_add_mismatched_window_dims_from_user_report2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_mismatched_window_dims_from_user_report2.onnx) ✅<br>[`scatter_add_mismatched_window_dims_from_user_report3`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_mismatched_window_dims_from_user_report3.onnx) ✅<br>[`scatter_add_fluids_pattern_updates_5_4_1_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_fluids_pattern_updates_5_4_1_1.onnx) ✅<br>[`scatter_add_in_cond_float64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_add_in_cond_float64.onnx) ✅ | v0.5.3 |
| [lax.scatter_mul](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scatter_mul.html) | [ScatterND](https://onnx.ai/onnx/operators/onnx__ScatterND.html) | [`scatter_mul_simple_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_simple_1d.onnx) ✅<br>[`scatter_mul_simple_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_simple_1d_f64.onnx) ✅<br>[`scatter_mul_window_2d_operand_1d_indices`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_window_2d_operand_1d_indices.onnx) ✅<br>[`scatter_mul_window_2d_operand_1d_indices_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_window_2d_operand_1d_indices_f64.onnx) ✅<br>[`scatter_mul_batch_updates_1d_operand`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_batch_updates_1d_operand.onnx) ✅<br>[`scatter_mul_batch_updates_1d_operand_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_batch_updates_1d_operand_f64.onnx) ✅<br>[`scatter_mul_mismatched_window_dims_from_user_report`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_mismatched_window_dims_from_user_report.onnx) ✅<br>[`scatter_mul_mismatched_window_dims_from_user_report2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_mismatched_window_dims_from_user_report2.onnx) ✅<br>[`scatter_mul_mismatched_window_dims_from_user_report3`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_mismatched_window_dims_from_user_report3.onnx) ✅<br>[`scatter_mul_fluids_pattern_updates_5_4_1_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_fluids_pattern_updates_5_4_1_1.onnx) ✅<br>[`scatter_mul_in_cond_float64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/scatter_mul_in_cond_float64.onnx) ✅ | v0.6.4 |
| [lax.select](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.select.html) | [Where](https://onnx.ai/onnx/operators/onnx__Where.html) | [`select_simple`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_simple.onnx) ✅<br>[`select_simple_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_simple_f64.onnx) ✅<br>[`select_mask_scores_tensor_else_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_mask_scores_tensor_else_dynamic.onnx) ✅<br>[`select_mask_scores_tensor_else_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_mask_scores_tensor_else_dynamic_f64.onnx) ✅<br>[`select_mask_scores_tensor_else`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_mask_scores_tensor_else.onnx) ✅<br>[`select_mask_scores_tensor_else_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_mask_scores_tensor_else_f64.onnx) ✅ | v0.7.1 |
| [lax.select_n](https://docs.jax.dev/en/latest/_autosummary/jax.lax.select_n.html) | [Where](https://onnx.ai/onnx/operators/onnx__Where.html) | [`select_n_bool_predicate_two_cases_float`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_bool_predicate_two_cases_float.onnx) ✅<br>[`select_n_bool_predicate_two_cases_float_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_bool_predicate_two_cases_float_f64.onnx) ✅<br>[`select_n_bool_predicate_two_cases_int`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_bool_predicate_two_cases_int.onnx) ✅<br>[`select_n_bool_predicate_two_cases_int_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_bool_predicate_two_cases_int_f64.onnx) ✅<br>[`select_n_bool_predicate_scalar_broadcast`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_bool_predicate_scalar_broadcast.onnx) ✅<br>[`select_n_bool_predicate_scalar_broadcast_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_bool_predicate_scalar_broadcast_f64.onnx) ✅<br>[`select_n_int_indices_three_cases`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_int_indices_three_cases.onnx) ✅<br>[`select_n_int_indices_three_cases_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_int_indices_three_cases_f64.onnx) ✅<br>[`select_n_int_indices_four_cases`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_int_indices_four_cases.onnx) ✅<br>[`select_n_int_indices_four_cases_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/select_n_int_indices_four_cases_f64.onnx) ✅ | v0.2.0 |
| [lax.sign](https://docs.jax.dev/en/latest/_autosummary/jax.lax.sign.html) | [Sign](https://onnx.ai/onnx/operators/onnx__Sign.html) | [`sign`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sign.onnx) ✅<br>[`sign_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sign_f64.onnx) ✅ | v0.5.0 |
| [lax.sin](https://docs.jax.dev/en/latest/_autosummary/jax.lax.sin.html) | [Sin](https://onnx.ai/onnx/operators/onnx__Sin.html) | [`sin`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sin.onnx) ✅<br>[`sin_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sin_f64.onnx) ✅ | v0.4.4 |
| [lax.sinh](https://docs.jax.dev/en/latest/_autosummary/jax.lax.sinh.html) | [Sinh](https://onnx.ai/onnx/operators/onnx__Sinh.html) | [`sinh`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sinh.onnx) ✅<br>[`sinh_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sinh_f64.onnx) ✅ | v0.4.4 |
| [lax.slice](https://docs.jax.dev/en/latest/_autosummary/jax.lax.slice.html) | [Slice](https://onnx.ai/onnx/operators/onnx__Slice.html) | [`slice_test1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/slice_test1.onnx) ✅<br>[`slice_test1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/slice_test1_f64.onnx) ✅<br>[`slice_3d_none_strides`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/slice_3d_none_strides.onnx) ✅<br>[`slice_3d_none_strides_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/slice_3d_none_strides_f64.onnx) ✅ | v0.1.0 |
| [lax.sort](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.sort.html) | [TopK](https://onnx.ai/onnx/operators/onnx__TopK.html) | [`sort_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sort_1d.onnx) ✅<br>[`sort_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sort_1d_f64.onnx) ✅<br>[`sort_2d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sort_2d.onnx) ✅<br>[`sort_2d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sort_2d_f64.onnx) ✅ | v0.2.0 |
| [lax.split](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.split.html) | [Split](https://onnx.ai/onnx/operators/onnx__Split.html) | [`lax_split_equal_parts`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_split_equal_parts.onnx) ✅<br>[`lax_split_equal_parts_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_split_equal_parts_f64.onnx) ✅<br>[`lax_split_unequal_parts`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_split_unequal_parts.onnx) ✅<br>[`lax_split_unequal_parts_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_split_unequal_parts_f64.onnx) ✅ | v0.7.2 |
| [lax.sqrt](https://docs.jax.dev/en/latest/_autosummary/jax.lax.sqrt.html) | [Sqrt](https://onnx.ai/onnx/operators/onnx__Sqrt.html) | [`sqrt`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sqrt.onnx) ✅<br>[`sqrt_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sqrt_f64.onnx) ✅ | v0.2.0 |
| [lax.square](https://docs.jax.dev/en/latest/_autosummary/jax.lax.square.html) | [Mul](https://onnx.ai/onnx/operators/onnx__Mul.html) | [`square`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/square.onnx) ✅<br>[`square_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/square_f64.onnx) ✅ | v0.2.0 |
| [lax.squeeze](https://docs.jax.dev/en/latest/_autosummary/jax.lax.squeeze.html) | [Squeeze](https://onnx.ai/onnx/operators/onnx__Squeeze.html) | [`lax_squeeze_specific_axis_0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_specific_axis_0.onnx) ✅<br>[`lax_squeeze_specific_axis_0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_specific_axis_0_f64.onnx) ✅<br>[`lax_squeeze_multiple_axes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_multiple_axes.onnx) ✅<br>[`lax_squeeze_multiple_axes_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_multiple_axes_f64.onnx) ✅<br>[`lax_squeeze_no_op_empty_dims`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_no_op_empty_dims.onnx) ✅<br>[`lax_squeeze_no_op_empty_dims_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_no_op_empty_dims_f64.onnx) ✅<br>[`lax_squeeze_problem_case_input_squeeze_only_axis_0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_problem_case_input_squeeze_only_axis_0.onnx) ✅<br>[`lax_squeeze_problem_case_input_squeeze_only_axis_0_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_problem_case_input_squeeze_only_axis_0_f64.onnx) ✅<br>[`lax_squeeze_problem_case_input_squeeze_axes_0_2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_problem_case_input_squeeze_axes_0_2.onnx) ✅<br>[`lax_squeeze_problem_case_input_squeeze_axes_0_2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_problem_case_input_squeeze_axes_0_2_f64.onnx) ✅<br>[`lax_squeeze_problem_case_input_squeeze_all_dims_explicitly`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_problem_case_input_squeeze_all_dims_explicitly.onnx) ✅<br>[`lax_squeeze_problem_case_input_squeeze_all_dims_explicitly_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/lax_squeeze_problem_case_input_squeeze_all_dims_explicitly_f64.onnx) ✅ | v0.2.0 |
| [lax.stop_gradient](https://docs.jax.dev/en/latest/_autosummary/jax.lax.stop_gradient.html) | [Identity](https://onnx.ai/onnx/operators/onnx__Identity.html) | [`stop_gradient`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/stop_gradient.onnx) ✅<br>[`stop_gradient_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/stop_gradient_f64.onnx) ✅ | v0.2.0 |
| [lax.sub](https://docs.jax.dev/en/latest/_autosummary/jax.lax.sub.html) | [Sub](https://onnx.ai/onnx/operators/onnx__Sub.html) | [`sub_test1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sub_test1.onnx) ✅<br>[`sub_test1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sub_test1_f64.onnx) ✅<br>[`sub_test2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sub_test2.onnx) ✅<br>[`sub_test2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/sub_test2_f64.onnx) ✅ | v0.1.0 |
| [lax.tanh](https://docs.jax.dev/en/latest/_autosummary/jax.lax.tanh.html) | [Tanh](https://onnx.ai/onnx/operators/onnx__Tanh.html) | [`tanh`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/tanh.onnx) ✅<br>[`tanh_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/tanh_f64.onnx) ✅ | v0.2.0 |
| [lax.transpose](https://docs.jax.dev/en/latest/_autosummary/jax.lax.transpose.html) | [Transpose](https://onnx.ai/onnx/operators/onnx__Transpose.html) | [`transpose_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/transpose_basic.onnx) ✅<br>[`transpose_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/transpose_basic_f64.onnx) ✅ | v0.2.0 |
| [lax.while_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html) | [Loop](https://onnx.ai/onnx/operators/onnx__Loop.html) | [`while_loop_counter`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_counter.onnx) ✅<br>[`while_loop_counter_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_counter_f64.onnx) ✅<br>[`while_loop_vector`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_vector.onnx) ✅<br>[`while_loop_vector_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_vector_f64.onnx) ✅<br>[`while_loop_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_f64.onnx) ✅<br>[`while_loop_multi_state_f32`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_multi_state_f32.onnx) ✅<br>[`while_loop_multi_state_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_multi_state_f64.onnx) ✅<br>[`while_loop_with_closure`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_with_closure.onnx) ✅<br>[`while_loop_with_closure_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_with_closure_f64.onnx) ✅<br>[`while_loop_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_basic.onnx) ✅<br>[`while_loop_two_state`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_two_state.onnx) ✅<br>[`while_loop_captured_tracer`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_captured_tracer.onnx) ✅<br>[`while_loop_with_scalar_state`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_with_scalar_state.onnx) ✅<br>[`while_loop_renamed_passthrough`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_renamed_passthrough.onnx) ✅<br>[`while_loop_closure_topo`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_closure_topo.onnx) ✅<br>[`while_loop_mixed_rank`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_mixed_rank.onnx) ✅<br>[`while_loop_tracer_passthrough`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_tracer_passthrough.onnx) ✅<br>[`while_loop_no_loop_output_reused_as_input`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_no_loop_output_reused_as_input.onnx) ✅<br>[`while_loop_4d_and_scalar_state`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_4d_and_scalar_state.onnx) ✅<br>[`while_loop_4d_and_scalar_state_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_4d_and_scalar_state_f64.onnx) ✅<br>[`while_loop_cnn_scalar_state_bug`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_cnn_scalar_state_bug.onnx) ✅<br>[`while_loop_cnn_scalar_state_bug_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_cnn_scalar_state_bug_f64.onnx) ✅<br>[`while_loop_nnx_repro`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_nnx_repro.onnx) ✅<br>[`while_loop_nnx_repro_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/lax/while_loop_nnx_repro_f64.onnx) ✅ | v0.5.1 |
| [nn.celu](https://docs.jax.dev/en/latest/_autosummary/jax.nn.celu.html) | [Celu](https://onnx.ai/onnx/operators/onnx__Celu.html) | [`jaxnn_celu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_celu.onnx) ✅<br>[`jaxnn_celu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_celu_f64.onnx) ✅<br>[`jaxnn_celu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_celu_1.onnx) ✅<br>[`jaxnn_celu_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_celu_1_f64.onnx) ✅ | v0.7.1 |
| [nn.dot_product_attention](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/attention.html#flax.nn.dot_product_attention) | [Add](https://onnx.ai/onnx/operators/onnx__Add.html)<br>[Cast](https://onnx.ai/onnx/operators/onnx__Cast.html)<br>[MatMul](https://onnx.ai/onnx/operators/onnx__MatMul.html)<br>[Mul](https://onnx.ai/onnx/operators/onnx__Mul.html)<br>[Not](https://onnx.ai/onnx/operators/onnx__Not.html)<br>[Softmax](https://onnx.ai/onnx/operators/onnx__Softmax.html)<br>[Transpose](https://onnx.ai/onnx/operators/onnx__Transpose.html)<br>[Where](https://onnx.ai/onnx/operators/onnx__Where.html) | [`dpa_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_basic.onnx) ✅<br>[`dpa_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_basic_f64.onnx) ✅<br>[`dpa_positional_bias_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_positional_bias_mask.onnx) ✅<br>[`dpa_positional_bias_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_positional_bias_mask_f64.onnx) ✅<br>[`dpa_diff_heads_embed`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_diff_heads_embed.onnx) ✅<br>[`dpa_diff_heads_embed_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_diff_heads_embed_f64.onnx) ✅<br>[`dpa_batch4_seq16`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_batch4_seq16.onnx) ✅<br>[`dpa_batch4_seq16_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_batch4_seq16_f64.onnx) ✅<br>[`dpa_float64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_float64.onnx) ✅<br>[`dpa_float64_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_float64_f64.onnx) ✅<br>[`dpa_heads1_embed4`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_heads1_embed4.onnx) ✅<br>[`dpa_heads1_embed4_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_heads1_embed4_f64.onnx) ✅<br>[`dpa_heads8_embed8`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_heads8_embed8.onnx) ✅<br>[`dpa_heads8_embed8_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_heads8_embed8_f64.onnx) ✅<br>[`dpa_batch1_seq2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_batch1_seq2.onnx) ✅<br>[`dpa_batch1_seq2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_batch1_seq2_f64.onnx) ✅<br>[`dpa_batch8_seq4`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_batch8_seq4.onnx) ✅<br>[`dpa_batch8_seq4_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_batch8_seq4_f64.onnx) ✅<br>[`dpa_axis1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_axis1.onnx) ✅<br>[`dpa_axis1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_axis1_f64.onnx) ✅<br>[`dpa_with_tensor_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_tensor_mask.onnx) ✅<br>[`dpa_with_tensor_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_tensor_mask_f64.onnx) ✅<br>[`dpa_tiny_mask_all_valid`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_tiny_mask_all_valid.onnx) ✅<br>[`dpa_tiny_mask_all_valid_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_tiny_mask_all_valid_f64.onnx) ✅<br>[`dpa_tiny_mask_mixed`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_tiny_mask_mixed.onnx) ✅<br>[`dpa_tiny_mask_mixed_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_tiny_mask_mixed_f64.onnx) ✅<br>[`dpa_one_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_one_false.onnx) ✅<br>[`dpa_one_false_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_one_false_f64.onnx) ✅<br>[`dpa_mostly_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_mostly_false.onnx) ✅<br>[`dpa_mostly_false_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_mostly_false_f64.onnx) ✅<br>[`dpa_with_causal_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_causal_mask.onnx) ✅<br>[`dpa_with_causal_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_causal_mask_f64.onnx) ✅<br>[`dpa_with_padding_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_padding_mask.onnx) ✅<br>[`dpa_with_padding_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_padding_mask_f64.onnx) ✅<br>[`dpa_with_local_window_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_local_window_mask.onnx) ✅<br>[`dpa_with_local_window_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_with_local_window_mask_f64.onnx) ✅<br>[`dpa_mask_none`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/dpa_mask_none.onnx) ✅ | v0.1.0 |
| [nn.elu](https://docs.jax.dev/en/latest/_autosummary/jax.nn.elu.html) | [Elu](https://onnx.ai/onnx/operators/onnx__Elu.html) | [`jaxnn_elu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_elu.onnx) ✅<br>[`jaxnn_elu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_elu_f64.onnx) ✅<br>[`jaxnn_elu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_elu_1.onnx) ✅<br>[`jaxnn_elu_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_elu_1_f64.onnx) ✅ | v0.7.1 |
| [nn.gelu](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.gelu.html) | [Gelu](https://onnx.ai/onnx/operators/onnx__Gelu.html) | [`jaxnn_gelu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_gelu.onnx) ✅<br>[`jaxnn_gelu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_gelu_f64.onnx) ✅<br>[`jaxnn_gelu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_gelu_1.onnx) ✅<br>[`jaxnn_gelu_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_gelu_1_f64.onnx) ✅<br>[`jaxnn_gelu_approx`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_gelu_approx.onnx) ✅<br>[`jaxnn_gelu_approx_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_gelu_approx_f64.onnx) ✅ | v0.7.1 |
| [nn.identity](https://docs.jax.dev/en/latest/_autosummary/jax.nn.identity.html) | [Identity](https://onnx.ai/onnx/operators/onnx__Identity.html) | [`jaxnn_identity`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_identity.onnx) ✅<br>[`jaxnn_identity_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_identity_f64.onnx) ✅<br>[`jaxnn_identity_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_identity_1.onnx) ✅<br>[`jaxnn_identity_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_identity_1_f64.onnx) ✅ | v0.7.1 |
| [nn.leaky_relu](https://docs.jax.dev/en/latest/_autosummary/jax.nn.leaky_relu.html) | [LeakyRelu](https://onnx.ai/onnx/operators/onnx__LeakyRelu.html) | [`jaxnn_leaky_relu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_leaky_relu.onnx) ✅<br>[`jaxnn_leaky_relu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_leaky_relu_f64.onnx) ✅<br>[`jaxnn_leaky_relu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_leaky_relu_1.onnx) ✅<br>[`jaxnn_leaky_relu_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_leaky_relu_1_f64.onnx) ✅ | v0.7.1 |
| [nn.mish](https://docs.jax.dev/en/latest/_autosummary/jax.nn.mish.html) | [Mish](https://onnx.ai/onnx/operators/onnx__Mish.html) | [`jaxnn_mish`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_mish.onnx) ✅<br>[`jaxnn_mish_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_mish_f64.onnx) ✅<br>[`jaxnn_mish_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_mish_1.onnx) ✅<br>[`jaxnn_mish_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_mish_1_f64.onnx) ✅ | v0.7.1 |
| [nn.relu](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.relu.html) | [Relu](https://onnx.ai/onnx/operators/onnx__Relu.html) | [`jaxnn_relu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_relu.onnx) ✅<br>[`jaxnn_relu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_relu_f64.onnx) ✅<br>[`jaxnn_relu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_relu_1.onnx) ✅<br>[`jaxnn_relu_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_relu_1_f64.onnx) ✅ | v0.7.1 |
| [nn.selu](https://docs.jax.dev/en/latest/_autosummary/jax.nn.selu.html) | [Selu](https://onnx.ai/onnx/operators/onnx__Selu.html) | [`jaxnn_selu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_selu.onnx) ✅<br>[`jaxnn_selu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_selu_f64.onnx) ✅<br>[`jaxnn_selu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_selu_1.onnx) ✅<br>[`jaxnn_selu_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_selu_1_f64.onnx) ✅ | v0.7.1 |
| [nn.sigmoid](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.sigmoid.html) | [Sigmoid](https://onnx.ai/onnx/operators/onnx__Sigmoid.html) | [`jaxnn_sigmoid`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_sigmoid.onnx) ✅<br>[`jaxnn_sigmoid_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_sigmoid_f64.onnx) ✅<br>[`jaxnn_sigmoid_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_sigmoid_1.onnx) ✅<br>[`jaxnn_sigmoid_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_sigmoid_1_f64.onnx) ✅ | v0.7.1 |
| [nn.soft_sign](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.soft_sign.html) | [Softsign](https://onnx.ai/onnx/operators/onnx__Softsign.html) | [`jaxnn_soft_sign`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_soft_sign.onnx) ✅<br>[`jaxnn_soft_sign_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_soft_sign_f64.onnx) ✅<br>[`jaxnn_soft_sign_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_soft_sign_1.onnx) ✅<br>[`jaxnn_soft_sign_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_soft_sign_1_f64.onnx) ✅ | v0.7.1 |
| [nn.softmax](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.softmax.html) | [Softmax](https://onnx.ai/onnx/operators/onnx__Softmax.html) | [`softmax`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/softmax.onnx) ✅<br>[`softmax_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/softmax_f64.onnx) ✅<br>[`softmax_2d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/softmax_2d.onnx) ✅<br>[`softmax_2d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/softmax_2d_f64.onnx) ✅<br>[`softmax_3d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/softmax_3d.onnx) ✅<br>[`softmax_3d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/softmax_3d_f64.onnx) ✅ | v0.1.0 |
| [nn.softplus](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.softplus.html) | [Softplus](https://onnx.ai/onnx/operators/onnx__Softplus.html) | [`jaxnn_softplus`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_softplus.onnx) ✅<br>[`jaxnn_softplus_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_softplus_f64.onnx) ✅<br>[`jaxnn_softplus_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_softplus_1.onnx) ✅<br>[`jaxnn_softplus_1_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/jaxnn_softplus_1_f64.onnx) ✅ | v0.7.1 |
| [nn.truncated_normal]() | ➖ | [`initializer`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/initializer.onnx) ✅<br>[`random_truncated_normal_positional`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/random_truncated_normal_positional.onnx) ✅<br>[`flax_dense_like_init`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nn/flax_dense_like_init.onnx) ✅ | v0.7.1 |
| [nnx.avg_pool](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.avg_pool) | [AveragePool](https://onnx.ai/onnx/operators/onnx__AveragePool.html)<br>[Transpose](https://onnx.ai/onnx/operators/onnx__Transpose.html) | [`avg_pool_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_dynamic.onnx) ✅<br>[`avg_pool`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool.onnx) ✅<br>[`avg_pool_same_padding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_same_padding_dynamic.onnx) ✅<br>[`avg_pool_same_padding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_same_padding.onnx) ✅<br>[`avg_pool_default_padding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_default_padding_dynamic.onnx) ✅<br>[`avg_pool_default_padding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_default_padding.onnx) ✅<br>[`avg_pool_stride1_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_stride1_dynamic.onnx) ✅<br>[`avg_pool_stride1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_stride1.onnx) ✅<br>[`avg_pool_win3x3_stride2_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_win3x3_stride2_dynamic.onnx) ✅<br>[`avg_pool_win3x3_stride2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_win3x3_stride2.onnx) ✅<br>[`avg_pool_stride_none_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_stride_none_dynamic.onnx) ✅<br>[`avg_pool_stride_none`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_stride_none.onnx) ✅<br>[`avg_pool_count_include_pad_false_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_count_include_pad_false_dynamic.onnx) ✅<br>[`avg_pool_count_include_pad_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/avg_pool_count_include_pad_false.onnx) ✅ | v0.1.0 |
| [nnx.batch_norm](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm) | [BatchNormalization](https://onnx.ai/onnx/operators/onnx__BatchNormalization.html) | [`batch_norm_no_bias_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_no_bias_no_scale_dynamic.onnx) ✅<br>[`batch_norm_no_bias_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_no_bias_no_scale.onnx) ✅<br>[`batch_norm_bias_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_bias_no_scale_dynamic.onnx) ✅<br>[`batch_norm_bias_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_bias_no_scale.onnx) ✅<br>[`batch_norm_no_bias_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_no_bias_scale_dynamic.onnx) ✅<br>[`batch_norm_no_bias_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_no_bias_scale.onnx) ✅<br>[`batch_norm_bias_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_bias_scale_dynamic.onnx) ✅<br>[`batch_norm_bias_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/batch_norm_bias_scale.onnx) ✅ | v0.1.0 |
| [nnx.conv](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.Conv) | [Conv](https://onnx.ai/onnx/operators/onnx__Conv.html)<br>[Transpose](https://onnx.ai/onnx/operators/onnx__Transpose.html) | [`conv_basic_bias_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_basic_bias_dynamic.onnx) ✅<br>[`conv_basic_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_basic_bias.onnx) ✅<br>[`conv_basic_bias_2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_basic_bias_2.onnx) ✅<br>[`conv_basic_bias_3`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_basic_bias_3.onnx) ✅<br>[`conv_stride2_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_stride2_bias.onnx) ✅<br>[`conv_no_bias_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_no_bias_dynamic.onnx) ✅<br>[`conv_no_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_no_bias.onnx) ✅<br>[`conv_valid_padding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_valid_padding.onnx) ✅<br>[`conv_stride1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_stride1.onnx) ✅<br>[`conv_stride2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_stride2.onnx) ✅<br>[`conv_different_kernel`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_different_kernel.onnx) ✅<br>[`conv_float64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_float64.onnx) ✅<br>[`conv_single_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_single_batch.onnx) ✅<br>[`conv_large_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/conv_large_batch.onnx) ✅ | v0.1.0 |
| [nnx.dot_product_attention](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/attention.html#flax.nnx.dot_product_attention) | [Cast](https://onnx.ai/onnx/operators/onnx__Cast.html)<br>[Div](https://onnx.ai/onnx/operators/onnx__Div.html)<br>[Einsum](https://onnx.ai/onnx/operators/onnx__Einsum.html)<br>[Gather](https://onnx.ai/onnx/operators/onnx__Gather.html)<br>[Shape](https://onnx.ai/onnx/operators/onnx__Shape.html)<br>[Softmax](https://onnx.ai/onnx/operators/onnx__Softmax.html)<br>[Sqrt](https://onnx.ai/onnx/operators/onnx__Sqrt.html) | [`dpa_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_basic.onnx) ✅<br>[`dpa_basic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_basic_f64.onnx) ✅<br>[`dpa_with_tensor_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_with_tensor_mask.onnx) ✅<br>[`dpa_with_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_with_bias.onnx) ✅<br>[`dpa_with_causal_mask`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_with_causal_mask.onnx) ✅<br>[`dpa_with_causal_mask_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_with_causal_mask_f64.onnx) ✅<br>[`dpa_with_mask_and_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_with_mask_and_bias.onnx) ✅<br>[`dpa_with_mask_and_bias_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dpa_with_mask_and_bias_f64.onnx) ✅ | v0.1.0 |
| [nnx.dropout](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/stochastic.html#flax.nnx.Dropout) | [Dropout](https://onnx.ai/onnx/operators/onnx__Dropout.html) | [`dropout_init_params_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_init_params_dynamic.onnx) ✅<br>[`dropout_init_params_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_init_params_dynamic_f64.onnx) ✅<br>[`dropout_init_params`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_init_params.onnx) ✅<br>[`dropout_init_params_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_init_params_f64.onnx) ✅<br>[`dropout_call_params_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_call_params_dynamic.onnx) ✅<br>[`dropout_call_params_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_call_params_dynamic_f64.onnx) ✅<br>[`dropout_call_params`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_call_params.onnx) ✅<br>[`dropout_call_params_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/dropout_call_params_f64.onnx) ✅ | v0.1.0 |
| [nnx.einsum](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.Einsum) | [Add](https://onnx.ai/onnx/operators/onnx__Add.html)<br>[Einsum](https://onnx.ai/onnx/operators/onnx__Einsum.html) | [`einsum_module_with_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/einsum_module_with_bias.onnx) ✅<br>[`einsum_module_with_bias_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/einsum_module_with_bias_f64.onnx) ✅<br>[`einsum_module_no_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/einsum_module_no_bias.onnx) ✅<br>[`einsum_module_no_bias_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/einsum_module_no_bias_f64.onnx) ✅ | v0.4.2 |
| [nnx.elu](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.elu.html) | [Elu](https://onnx.ai/onnx/operators/onnx__Elu.html) | [`elu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/elu.onnx) ✅ | v0.1.0 |
| [nnx.embed](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.Embed) | [Gather](https://onnx.ai/onnx/operators/onnx__Gather.html) | [`token_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/token_embedding_dynamic.onnx) ✅<br>[`token_embedding_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/token_embedding_dynamic_f64.onnx) ✅<br>[`token_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/token_embedding.onnx) ✅<br>[`token_embedding_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/token_embedding_f64.onnx) ✅<br>[`positional_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/positional_embedding_dynamic.onnx) ✅<br>[`positional_embedding_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/positional_embedding_dynamic_f64.onnx) ✅<br>[`positional_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/positional_embedding.onnx) ✅<br>[`positional_embedding_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/positional_embedding_f64.onnx) ✅ | v0.7.0 |
| [nnx.gelu](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/activations.html#flax.nnx.gelu) | [Gelu](https://onnx.ai/onnx/operators/onnx__Gelu.html) | [`gelu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu.onnx) ✅<br>[`gelu_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_1.onnx) ✅<br>[`gelu_2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_2.onnx) ✅<br>[`gelu_2_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_2_f64.onnx) ✅<br>[`gelu_3_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_3_dynamic.onnx) ✅<br>[`gelu_3_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_3_dynamic_f64.onnx) ✅<br>[`gelu_3`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_3.onnx) ✅<br>[`gelu_3_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/gelu_3_f64.onnx) ✅ | v0.1.0 |
| [nnx.group_norm](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.GroupNorm) | [GroupNormalization](https://onnx.ai/onnx/operators/onnx__GroupNormalization.html) | [`group_norm`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm.onnx) ✅<br>[`group_norm_no_bias_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_no_bias_no_scale_dynamic.onnx) ✅<br>[`group_norm_no_bias_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_no_bias_no_scale.onnx) ✅<br>[`group_norm_bias_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_bias_no_scale_dynamic.onnx) ✅<br>[`group_norm_bias_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_bias_no_scale.onnx) ✅<br>[`group_norm_no_bias_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_no_bias_scale_dynamic.onnx) ✅<br>[`group_norm_no_bias_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_no_bias_scale.onnx) ✅<br>[`group_norm_bias_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_bias_scale_dynamic.onnx) ✅<br>[`group_norm_bias_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/group_norm_bias_scale.onnx) ✅ | v0.3.0 |
| [nnx.layer_norm](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.LayerNorm) | [LayerNormalization](https://onnx.ai/onnx/operators/onnx__LayerNormalization.html) | [`layer_norm_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_dynamic.onnx) ✅<br>[`layer_norm`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm.onnx) ✅<br>[`layer_norm_no_bias_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_no_bias_no_scale_dynamic.onnx) ✅<br>[`layer_norm_no_bias_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_no_bias_no_scale.onnx) ✅<br>[`layer_norm_bias_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_bias_no_scale_dynamic.onnx) ✅<br>[`layer_norm_bias_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_bias_no_scale.onnx) ✅<br>[`layer_norm_no_bias_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_no_bias_scale_dynamic.onnx) ✅<br>[`layer_norm_no_bias_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_no_bias_scale.onnx) ✅<br>[`layer_norm_bias_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_bias_scale_dynamic.onnx) ✅<br>[`layer_norm_bias_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_bias_scale.onnx) ✅<br>[`layer_norm_multiaxis_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_multiaxis_dynamic.onnx) ✅<br>[`layer_norm_multiaxis`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_multiaxis.onnx) ✅<br>[`layer_norm_symbolic_batch_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_symbolic_batch_dynamic.onnx) ✅<br>[`layer_norm_symbolic_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/layer_norm_symbolic_batch.onnx) ✅ | v0.1.0 |
| [nnx.leaky_relu](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.leaky_relu.html) | [LeakyRelu](https://onnx.ai/onnx/operators/onnx__LeakyRelu.html) | [`leaky_relu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/leaky_relu.onnx) ✅ | v0.1.0 |
| [nnx.linear](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html) | [Gemm](https://onnx.ai/onnx/operators/onnx__Gemm.html)<br>[Reshape](https://onnx.ai/onnx/operators/onnx__Reshape.html) | [`linear_symbolic_batch_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_symbolic_batch_dynamic.onnx) ✅<br>[`linear_symbolic_batch_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_symbolic_batch_dynamic_f64.onnx) ✅<br>[`linear_symbolic_batch`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_symbolic_batch.onnx) ✅<br>[`linear_symbolic_batch_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_symbolic_batch_f64.onnx) ✅<br>[`linear_high_rank_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_dynamic.onnx) ✅<br>[`linear_high_rank_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_dynamic_f64.onnx) ✅<br>[`linear_high_rank`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank.onnx) ✅<br>[`linear_high_rank_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_f64.onnx) ✅<br>[`linear_no_bias_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_no_bias_dynamic.onnx) ✅<br>[`linear_no_bias_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_no_bias_dynamic_f64.onnx) ✅<br>[`linear_no_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_no_bias.onnx) ✅<br>[`linear_no_bias_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_no_bias_f64.onnx) ✅<br>[`linear_high_rank_no_bias_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_no_bias_dynamic.onnx) ✅<br>[`linear_high_rank_no_bias_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_no_bias_dynamic_f64.onnx) ✅<br>[`linear_high_rank_no_bias`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_no_bias.onnx) ✅<br>[`linear_high_rank_no_bias_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_high_rank_no_bias_f64.onnx) ✅ | v0.1.0 |
| [nnx.linear_general](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.LinearGeneral) | [Gemm](https://onnx.ai/onnx/operators/onnx__Gemm.html)<br>[Reshape](https://onnx.ai/onnx/operators/onnx__Reshape.html) | [`linear_general_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general_dynamic.onnx) ✅<br>[`linear_general`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general.onnx) ✅<br>[`linear_general_2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general_2.onnx) ✅<br>[`linear_general_3`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general_3.onnx) ✅<br>[`linear_general_4`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general_4.onnx) ✅<br>[`linear_general_abstract_eval_axes`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general_abstract_eval_axes.onnx) ✅<br>[`linear_general_abstract_eval_axes_pair`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/linear_general_abstract_eval_axes_pair.onnx) ✅ | v0.1.0 |
| [nnx.log_softmax](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.log_softmax.html) | [LogSoftmax](https://onnx.ai/onnx/operators/onnx__LogSoftmax.html) | [`log_softmax`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/log_softmax.onnx) ✅<br>[`log_softmax_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/log_softmax_f64.onnx) ✅ | v0.1.0 |
| [nnx.max_pool](https://flax-linen.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.max_pool) | [MaxPool](https://onnx.ai/onnx/operators/onnx__MaxPool.html)<br>[Transpose](https://onnx.ai/onnx/operators/onnx__Transpose.html) | [`max_pool`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/max_pool.onnx) ✅<br>[`max_pool_same_padding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/max_pool_same_padding.onnx) ✅ | v0.1.0 |
| [nnx.relu](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.relu.html) | [Relu](https://onnx.ai/onnx/operators/onnx__Relu.html) | [`relu_1d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/relu_1d.onnx) ✅<br>[`relu_1d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/relu_1d_f64.onnx) ✅<br>[`relu_4d_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/relu_4d_dynamic.onnx) ✅<br>[`relu_4d_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/relu_4d_dynamic_f64.onnx) ✅<br>[`relu_4d`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/relu_4d.onnx) ✅<br>[`relu_4d_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/relu_4d_f64.onnx) ✅ | v0.1.0 |
| [nnx.rms_norm](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.RMSNorm) | [RMSNormalization](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html) | [`rms_norm_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/rms_norm_basic.onnx) ✅<br>[`rms_norm_use_scale_false`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/rms_norm_use_scale_false.onnx) ✅<br>[`rms_norm_4d_dynamic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/rms_norm_4d_dynamic_dynamic.onnx) ✅<br>[`rms_norm_4d_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/rms_norm_4d_dynamic.onnx) ✅<br>[`rms_norm_4d_dynamic_no_scale_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/rms_norm_4d_dynamic_no_scale_dynamic.onnx) ✅<br>[`rms_norm_4d_dynamic_no_scale`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/rms_norm_4d_dynamic_no_scale.onnx) ✅ | v0.3.0 |
| [nnx.sigmoid](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/activations.html#flax.nnx.sigmoid) | [Sigmoid](https://onnx.ai/onnx/operators/onnx__Sigmoid.html) | [`sigmoid`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/sigmoid.onnx) ✅<br>[`sigmoid_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/sigmoid_f64.onnx) ✅ | v0.1.0 |
| [nnx.softmax](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.softmax.html) | [Softmax](https://onnx.ai/onnx/operators/onnx__Softmax.html) | [`softmax_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/softmax_dynamic.onnx) ✅<br>[`softmax_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/softmax_dynamic_f64.onnx) ✅<br>[`softmax`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/softmax.onnx) ✅<br>[`softmax_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/softmax_f64.onnx) ✅ | v0.1.0 |
| [nnx.softplus](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.softplus.html) | [Softplus](https://onnx.ai/onnx/operators/onnx__Softplus.html) | [`softplus`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/softplus.onnx) ✅ | v0.1.0 |
| [nnx.tanh](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/activations.html#flax.nnx.tanh) | [Tanh](https://onnx.ai/onnx/operators/onnx__Tanh.html) | [`tanh`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/tanh.onnx) ✅<br>[`tanh_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/primitives/nnx/tanh_f64.onnx) ✅ | v0.1.0 |

<!-- AUTOGENERATED TABLE END -->

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

---

## 🎯 Examples

<!-- AUTOGENERATED EXAMPLES TABLE START -->

| Component | Description | Testcases | Since |
|:----------|:------------|:----------|:------|
| SimpleLinearExample | A simple linear layer example using Equinox. | [`simple_linear_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/simple_linear_dynamic.onnx) ✅<br>[`simple_linear_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/simple_linear_dynamic_f64.onnx) ✅<br>[`simple_linear`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/simple_linear.onnx) ✅<br>[`simple_linear_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/simple_linear_f64.onnx) ✅<br>[`nn_linear_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/nn_linear_dynamic.onnx) ✅<br>[`nn_linear_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/nn_linear_dynamic_f64.onnx) ✅<br>[`nn_linear`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/nn_linear.onnx) ✅<br>[`nn_linear_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/eqx/nn_linear_f64.onnx) ✅ | v0.7.1 |
| GPT | A simple GPT model that reuses nnx.MultiHeadAttention. | [`gpt_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_dynamic.onnx) ✅<br>[`gpt`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt.onnx) ✅ | v0.7.0 |
| GPT_Attention | A multi-head attention layer. | [`gpt_attention`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_attention.onnx) ✅ | v0.7.1 |
| GPT_CausalSelfAttention | A causal self-attention module. | [`causal_self_attention_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/causal_self_attention_dynamic.onnx) ✅<br>[`causal_self_attention`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/causal_self_attention.onnx) ✅ | v0.7.0 |
| GPT_Embeddings | Combines token and position embeddings with dropout. | [`gpt_embeddings_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_embeddings_dynamic.onnx) ✅<br>[`gpt_embeddings`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_embeddings.onnx) ✅ | v0.7.0 |
| GPT_Head | The head of the GPT model. | [`gpt_head_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_head_dynamic.onnx) ✅<br>[`gpt_head`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_head.onnx) ✅ | v0.7.0 |
| GPT_MLP | An MLP block with GELU activation from nanoGPT. | [`gpt_mlp_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_mlp_dynamic.onnx) ✅<br>[`gpt_mlp`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_mlp.onnx) ✅ | v0.7.0 |
| GPT_PositionEmbedding | A positional embedding layer using nnx.Embed. | [`position_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/position_embedding.onnx) ✅ | v0.7.0 |
| GPT_TokenEmbedding | A token embedding layer using nnx.Embed. | [`token_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/token_embedding_dynamic.onnx) ✅<br>[`token_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/token_embedding.onnx) ✅ | v0.7.0 |
| GPT_TransformerBlock | A transformer block combining attention and MLP. | [`gpt_block_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_block_dynamic.onnx) ✅<br>[`gpt_block`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/gpt_block.onnx) ✅ | v0.7.0 |
| GPT_TransformerStack | A stack of transformer blocks. | [`transformer_stack_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/transformer_stack_dynamic.onnx) ✅<br>[`transformer_stack`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/transformer_stack.onnx) ✅ | v0.7.0 |
| broadcast_add | Simple dynamic broadcast + add | [`broadcast_add_dynamic_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/broadcast_add_dynamic_dynamic.onnx) ✅<br>[`broadcast_add_dynamic_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/broadcast_add_dynamic_dynamic_f64.onnx) ✅<br>[`broadcast_add_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/broadcast_add_dynamic.onnx) ✅<br>[`broadcast_add_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/gpt/broadcast_add_dynamic_f64.onnx) ✅ | v0.7.0 |
| cfl_timestep | Tests the CFL condition timestep calculation. | [`cfl_timestep_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jaxfluids/cfl_timestep_f64.onnx) ✅ | v0.6.5 |
| weno_reconstruction | Tests the complex arithmetic pattern found in WENO schemes. | [`weno_reconstruction_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jaxfluids/weno_reconstruction_f64.onnx) ✅ | v0.6.5 |
| fori_loop_test | fori_loop_test: Demonstrates jax.lax.fori_loop with a simple loop. | [`fori_loop_test`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/fori_loop_test.onnx) ✅<br>[`fori_loop_test_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/fori_loop_test_f64.onnx) ➖ | v0.6.3 |
| issue18_abs | Test jnp.abs from issue 18 | [`abs_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/abs_fn.onnx) ✅<br>[`abs_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/abs_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_arange | Test arange from issue 18 | [`arange_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/arange_fn.onnx) ✅<br>[`arange_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/arange_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_fori_loop | Test fori_loop from issue 18 | [`fori_loop_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/fori_loop_fn.onnx) ✅<br>[`fori_loop_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/fori_loop_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_linspace | Test linspace from issue 18 | [`linspace_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/linspace_fn.onnx) ✅<br>[`linspace_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/linspace_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_scan | Test scan from issue 18 (no xs) | [`scan_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/scan_fn.onnx) ✅<br>[`scan_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/scan_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_sign | Test jnp.sign from issue 18 | [`sign_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/sign_fn.onnx) ✅<br>[`sign_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/sign_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_where | Test where from issue 18 | [`where_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/where_fn.onnx) ✅<br>[`where_fn_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/where_fn_f64.onnx) ✅ | v0.6.3 |
| issue18_while_loop | Test while_loop from issue 18 | [`while_loop_fn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/while_loop_fn.onnx) ✅ | v0.6.3 |
| select_test | select_test: Demonstrates jnp.select with a dynamic condition based on an input array. | [`select_test_all_options`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/select_test_all_options.onnx) ➖<br>[`select_test_scalar_select_option_0`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/select_test_scalar_select_option_0.onnx) ➖<br>[`select_test_scalar_select_option_1`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/select_test_scalar_select_option_1.onnx) ➖<br>[`select_test_scalar_select_option_2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/select_test_scalar_select_option_2.onnx) ➖<br>[`select_test_default_case`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/select_test_default_case.onnx) ➖ | v0.6.1 |
| sort_test | sort_test: Demonstrates jnp.sort on slices of an input array. | [`sort_test_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/jnp/sort_test_basic.onnx) ➖ | v0.6.1 |
| cond_scatter_add_mul | Tests scatter_add/mul inside jnp.where branches | [`cond_scatter_add_mul_f64_a`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/lax/cond_scatter_add_mul_f64_a.onnx) ✅<br>[`cond_scatter_add_mul_f64_b`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/lax/cond_scatter_add_mul_f64_b.onnx) ✅ | v0.6.4 |
| cond_scatter_repro | Reproduces a bug where lax.cond subgraphs do not inherit parent initializers. | [`cond_scatter_repro_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/lax/cond_scatter_repro_f64.onnx) ✅ | v0.6.4 |
| remat2 | Tests a simple case of `jax.checkpoint` (also known as `jax.remat2`). | [`checkpoint_scalar_f32`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/lax/checkpoint_scalar_f32.onnx) ✅<br>[`checkpoint_scalar_f32_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/lax/checkpoint_scalar_f32_f64.onnx) ✅ | v0.6.5 |
| AutoEncoder | A simple autoencoder example. | [`simple_autoencoder`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_autoencoder.onnx) ✅<br>[`simple_autoencoder_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_autoencoder_f64.onnx) ✅ | v0.2.0 |
| CNN | A simple convolutional neural network (CNN). | [`simple_cnn_explicit_dimensions`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_cnn_explicit_dimensions.onnx) ✅<br>[`simple_cnn_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_cnn_dynamic.onnx) ✅<br>[`simple_cnn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_cnn.onnx) ✅ | v0.1.0 |
| ForiLoop | fori_loop example | [`fori_loop_counter`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/fori_loop_counter.onnx) ✅<br>[`fori_loop_counter_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/fori_loop_counter_f64.onnx) ✅ | v0.5.1 |
| GRUCell | Vanilla gated-recurrent-unit cell from **Flax/nnx**. There is no 1-to-1 ONNX operator, so the converter decomposes it into MatMul, Add, Sigmoid, Tanh, etc. | [`gru_cell_basic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/gru_cell_basic.onnx) ✅ | v0.7.2 |
| MLP | A simple Multi-Layer Perceptron (MLP) with BatchNorm, Dropout, and GELU activation. | [`simple_mlp_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_dynamic.onnx) ✅<br>[`simple_mlp_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_dynamic_f64.onnx) ✅<br>[`simple_mlp`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp.onnx) ✅<br>[`simple_mlp_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_f64.onnx) ✅<br>[`simple_mlp_with_call_params_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_with_call_params_dynamic.onnx) ✅<br>[`simple_mlp_with_call_params_dynamic_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_with_call_params_dynamic_f64.onnx) ✅<br>[`simple_mlp_with_call_params`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_with_call_params.onnx) ✅<br>[`simple_mlp_with_call_params_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/simple_mlp_with_call_params_f64.onnx) ✅ | 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_nn_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/multihead_attention_nn_dynamic.onnx) ✅<br>[`multihead_attention_nn`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/multihead_attention_nn.onnx) ✅<br>[`multihead_attention_nnx_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/multihead_attention_nnx_dynamic.onnx) ✅<br>[`multihead_attention_nnx`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/multihead_attention_nnx.onnx) ✅ | v0.2.0 |
| SequentialReLU | Two ReLU activations chained with nnx.Sequential (no parameters). | [`sequential_double_relu`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/sequential_double_relu.onnx) ✅<br>[`sequential_double_relu_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/sequential_double_relu_f64.onnx) ✅ | v0.7.1 |
| SequentialWithResidual | Tests nnx.Sequential nested inside a module with a residual connection. | [`sequential_nested_with_residual`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/sequential_nested_with_residual.onnx) ✅<br>[`sequential_nested_with_residual_f64`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/sequential_nested_with_residual_f64.onnx) ✅ | v0.7.1 |
| TransformerDecoder | A single-layer Transformer decoder built with nnx primitives (MHA, LayerNorm, Feed-Forward, Dropout). | [`tiny_decoder`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/tiny_decoder.onnx) ✅ | v0.7.1 |
| TransformerDecoder2 | A single-layer Transformer decoder built with nnx primitives (MHA, LayerNorm, Feed-Forward, Dropout). | [`tiny_decoder2`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/nnx/tiny_decoder2.onnx) ✅ | v0.7.1 |
| onnx_functions_000 | one function on an outer layer. | [`000_one_function_on_outer_layer_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/000_one_function_on_outer_layer_dynamic.onnx) ✅<br>[`000_one_function_on_outer_layer`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/000_one_function_on_outer_layer.onnx) ✅ | v0.4.0 |
| onnx_functions_001 | one function on an inner layer. | [`001_one_function_inner_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/001_one_function_inner_dynamic.onnx) ✅<br>[`001_one_function_inner`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/001_one_function_inner.onnx) ✅ | v0.4.0 |
| onnx_functions_002 | two nested functions. | [`002_two_nested_functions_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/002_two_nested_functions_dynamic.onnx) ✅<br>[`002_two_nested_functions`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/002_two_nested_functions.onnx) ✅ | v0.4.0 |
| onnx_functions_003 | two nested functions. | [`003_two_simple_nested_functions_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/003_two_simple_nested_functions_dynamic.onnx) ✅<br>[`003_two_simple_nested_functions`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/003_two_simple_nested_functions.onnx) ✅ | v0.4.0 |
| onnx_functions_004 | nested function plus component | [`004_nested_function_plus_component_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/004_nested_function_plus_component_dynamic.onnx) ✅<br>[`004_nested_function_plus_component`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/004_nested_function_plus_component.onnx) ✅ | v0.4.0 |
| onnx_functions_005 | nested function plus more components | [`005_nested_function_plus_component_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/005_nested_function_plus_component_dynamic.onnx) ✅<br>[`005_nested_function_plus_component`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/005_nested_function_plus_component.onnx) ✅ | v0.4.0 |
| onnx_functions_006 | one function on an outer layer. | [`006_one_function_outer_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/006_one_function_outer_dynamic.onnx) ✅<br>[`006_one_function_outer`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/006_one_function_outer.onnx) ✅ | v0.4.0 |
| onnx_functions_007 | transformer block with nested mlp block with call parameter | [`007_transformer_block_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/007_transformer_block_dynamic.onnx) ✅<br>[`007_transformer_block`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/007_transformer_block.onnx) ✅ | v0.4.0 |
| onnx_functions_008 | transformer block with nested mlp block no call parameter | [`008_transformer_block_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/008_transformer_block_dynamic.onnx) ✅<br>[`008_transformer_block`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/008_transformer_block.onnx) ✅ | v0.4.0 |
| onnx_functions_009 | transformer block using decorator on class and function | [`009_transformer_block_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/009_transformer_block_dynamic.onnx) ✅<br>[`009_transformer_block`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/009_transformer_block.onnx) ✅ | v0.4.0 |
| onnx_functions_010 | transformer stack | [`010_transformer_stack_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/010_transformer_stack_dynamic.onnx) ✅<br>[`010_transformer_stack`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/010_transformer_stack.onnx) ✅ | v0.4.0 |
| onnx_functions_012 | Vision Transformer (ViT) | [`012_vit_conv_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/012_vit_conv_embedding_dynamic.onnx) ✅<br>[`012_vit_conv_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/012_vit_conv_embedding.onnx) ✅ | v0.4.0 |
| onnx_functions_013 | Vision Transformer (ViT) | [`013_vit_conv_embedding_with_call_params_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/013_vit_conv_embedding_with_call_params_dynamic.onnx) ✅<br>[`013_vit_conv_embedding_with_call_params`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/013_vit_conv_embedding_with_call_params.onnx) ✅<br>[`013_vit_conv_embedding_with_internal_call_params_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/013_vit_conv_embedding_with_internal_call_params_dynamic.onnx) ✅<br>[`013_vit_conv_embedding_with_internal_call_params`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/013_vit_conv_embedding_with_internal_call_params.onnx) ✅ | v0.4.0 |
| onnx_functions_014 | one function on an outer layer. | [`014_one_function_with_input_param_with_default_value`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/014_one_function_with_input_param_with_default_value.onnx) ✅<br>[`014_one_function_without_input_param_with_default_value_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/014_one_function_without_input_param_with_default_value_dynamic.onnx) ✅<br>[`014_one_function_without_input_param_with_default_value`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/014_one_function_without_input_param_with_default_value.onnx) ✅ | v0.4.0 |
| onnx_functions_015 | one function on an outer layer. | [`015_one_function_with_input_param_without_default_value_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/015_one_function_with_input_param_without_default_value_dynamic.onnx) ✅<br>[`015_one_function_with_input_param_without_default_value`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/015_one_function_with_input_param_without_default_value.onnx) ✅ | v0.4.0 |
| onnx_functions_016 | nested function plus more components | [`016_internal_function_with_input_param_with_default_value_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/016_internal_function_with_input_param_with_default_value_dynamic.onnx) ✅<br>[`016_internal_function_with_input_param_with_default_value`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/onnx_functions/016_internal_function_with_input_param_with_default_value.onnx) ✅ | v0.4.0 |
| ClassificationHead | Classification head for Vision Transformer | [`classification_head_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/classification_head_dynamic.onnx) ✅<br>[`classification_head`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/classification_head.onnx) ✅ | v0.4.0 |
| ClassificationHeadFlatten | Classification head for Vision Transformer | [`classification_head_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/classification_head_flat_dynamic.onnx) ✅<br>[`classification_head_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/classification_head_flat.onnx) ✅ | v0.4.0 |
| ConcatClsToken | Concatenate CLS token to the input embedding | [`concat_cls_token_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/concat_cls_token_dynamic.onnx) ✅<br>[`concat_cls_token`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/concat_cls_token.onnx) ✅ | v0.4.0 |
| ConcatClsTokenFlatten | Concatenate CLS token to the input embedding | [`concat_cls_token_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/concat_cls_token_flat_dynamic.onnx) ✅<br>[`concat_cls_token_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/concat_cls_token_flat.onnx) ✅ | v0.4.0 |
| ConvEmbedding | Convolutional Token Embedding for MNIST with hierarchical downsampling. | [`mnist_conv_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/mnist_conv_embedding_dynamic.onnx) ✅<br>[`mnist_conv_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/mnist_conv_embedding.onnx) ✅ | v0.1.0 |
| ConvEmbeddingFlatten | Convolutional Token Embedding for MNIST with hierarchical downsampling. | [`mnist_conv_embedding_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/mnist_conv_embedding_flat_dynamic.onnx) ✅<br>[`mnist_conv_embedding_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/mnist_conv_embedding_flat.onnx) ✅ | v0.1.0 |
| FeedForward | MLP in Transformer | [`feed_forward_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/feed_forward_dynamic.onnx) ✅<br>[`feed_forward`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/feed_forward.onnx) ✅ | v0.1.0 |
| FeedForwardFlatten | MLP in Transformer | [`feed_forward_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/feed_forward_flat_dynamic.onnx) ✅<br>[`feed_forward_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/feed_forward_flat.onnx) ✅ | v0.1.0 |
| GetToken | Get the CLS token from the input embedding | [`get_token_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/get_token_dynamic.onnx) ✅<br>[`get_token`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/get_token.onnx) ✅ | v0.4.0 |
| GetTokenFlatten | Get the CLS token from the input embedding | [`get_token_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/get_token_flat_dynamic.onnx) ✅<br>[`get_token_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/get_token_flat.onnx) ✅ | v0.4.0 |
| PatchEmbedding | Cutting the image into patches and linearly embedding them. | [`patch_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/patch_embedding_dynamic.onnx) ✅<br>[`patch_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/patch_embedding.onnx) ✅ | v0.1.0 |
| PatchEmbeddingFlatten | Cutting the image into patches and linearly embedding them. | [`patch_embedding_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/patch_embedding_flat_dynamic.onnx) ✅<br>[`patch_embedding_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/patch_embedding_flat.onnx) ✅ | v0.1.0 |
| PositionalEmbedding | Add positional embedding to the input embedding | [`positional_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/positional_embedding_dynamic.onnx) ✅<br>[`positional_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/positional_embedding.onnx) ✅ | v0.4.0 |
| PositionalEmbeddingFlatten | Add positional embedding to the input embedding | [`positional_embedding_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/positional_embedding_flat_dynamic.onnx) ✅<br>[`positional_embedding_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/positional_embedding_flat.onnx) ✅ | v0.4.0 |
| TransformerBlock | Transformer from 'Attention Is All You Need.' | [`transformer_block_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_block_dynamic.onnx) ✅<br>[`transformer_block`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_block.onnx) ✅ | v0.1.0 |
| TransformerBlockFlatten | Transformer from 'Attention Is All You Need.' | [`transformer_block_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_block_flat_dynamic.onnx) ✅<br>[`transformer_block_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_block_flat.onnx) ✅ | v0.1.0 |
| TransformerStack | Stack of Transformer blocks | [`transformer_stack_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_stack_dynamic.onnx) ✅<br>[`transformer_stack`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_stack.onnx) ✅ | v0.1.0 |
| TransformerStackFlatten | Stack of Transformer blocks | [`transformer_stack_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_stack_flat_dynamic.onnx) ✅<br>[`transformer_stack_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/transformer_stack_flat.onnx) ✅ | v0.1.0 |
| VisionTransformer | A Vision Transformer (ViT) model for MNIST with configurable embedding type. | [`vit_conv_embedding_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_conv_embedding_dynamic.onnx) ✅<br>[`vit_conv_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_conv_embedding.onnx) ✅<br>[`vit_patch_embedding`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_patch_embedding.onnx) ✅ | v0.2.0 |
| VisionTransformerFlatten | A Vision Transformer (ViT) model for MNIST with configurable embedding type. | [`vit_conv_embedding_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_conv_embedding_flat_dynamic.onnx) ✅<br>[`vit_conv_embedding_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_conv_embedding_flat.onnx) ✅<br>[`vit_patch_embedding_flat_dynamic`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_patch_embedding_flat_dynamic.onnx) ✅<br>[`vit_patch_embedding_flat`](https://netron.app/?url=https://huggingface.co/enpasos/jax2onnx-models/resolve/main/examples/vit/vit_patch_embedding_flat.onnx) ✅ | v0.2.0 |

<!-- AUTOGENERATED EXAMPLES TABLE END -->

---

## 📌 Dependencies

**Versions of Major Dependencies:**

| Library       | Versions |  
|:--------------|:---------| 
| `JAX`         | 0.6.2    | 
| `Flax`        | 0.10.7   | 
| `onnx`        | 1.18.0   |  
| `onnxruntime` | 1.22.1   |  

*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`](./jax2onnx/plugins). 
a custom primitive or an example.
- **Bug fixes & improvements:** PRs and issues are always welcome.

---

## 💾 Installation

Install from PyPI:

```bash
pip install jax2onnx  
```


---

## 📜 License

This project is licensed under the Apache License, Version 2.0. See [`LICENSE`](./LICENSE) for details.

---

## 🌟 Special Thanks

Special thanks for example contributions to [@burakssen](https://github.com/burakssen), [@Cadynum](https://github.com/Cadynum) and [@clementpoiret](https://github.com/clementpoiret)

Special thanks for plugin contributions to [@burakssen](https://github.com/burakssen) and [@clementpoiret](https://github.com/clementpoiret)

Special thanks to [tumaer/JAXFLUIDS](https://github.com/tumaer/JAXFLUIDS) for contributing valuable insights rooted in physics simulation use cases.

Special thanks to [@lutzroeder](https://github.com/lutzroeder) for making shapes internal to ONNX function visible in his great Netron viewer.

- [ONNX: Function value_info support #1447](https://github.com/lutzroeder/netron/issues/1447)


Special thanks to the community members involved in:

- [Flax Feature Request #4430](https://github.com/google/flax/issues/4430)
- [JAX Feature Request #26430](https://github.com/jax-ml/jax/issues/26430)

A huge thanks especially to [@limarta](https://github.com/limarta), whose elegant [jaxpr-to-ONNX demonstration](https://gist.github.com/limarta/855a88cc1c0163487a9dc369891147ab) significantly inspired this project.

---

**Happy converting! 🎉**



