Skip to content

ONNX Operator Coverage

This matrix starts from the official ONNX operator index and shows which operators are referenced in jax2onnx/plugins. It is a public reference for coverage analysis; Next Action classifies gaps and metadata/lowering drift, but it is not a per-release commitment list.

  • Metadata: an onnx__<Op>.html operator link was found in plugin metadata, including links generated from literal operator lists.
  • Lowering: plugin code calls builder.<Op>(...), builder.op(...), or a shared unary/reduction emitter with a statically known operator name.
  • Potential JAX Ops: candidate JAX APIs that naturally map to the ONNX operator.
  • Source index: ONNX 1.24.0 documentation (https://onnx.ai/onnx/operators/index.html)
  • Operators in index: 204
  • Operators referenced in plugins: 155
  • Coverage: 76.0%
  • Open operators by next action:
  • Sequence/Optional/String containers: 16
  • Quantization scope: 6
  • Vision-specific native ops: 9
  • Target-model specialized ops: 12
  • Low-complexity math quick wins: 0
  • Common model native ops: 0
  • General triage: 6
  • Potential JAX Ops lists candidate JAX entry points for each operator.
  • Next Action classifies uncovered operators and metadata/lowering drift.
  • Extra plugin names not present in index: 6
ONNX Operator In Plugins Metadata Lowering Plugin Modules Potential JAX Ops Next Action
Abs ✅ ✅ ✅ jax/lax/abs
jax/lax/bessel_i0e
jax/lax/bessel_i1e
jax/lax/cbrt
jax/lax/lgamma
jax/lax/lu
... +11
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Acos ✅ ✅ ✅ jax/lax/acos
jax/numpy/acos
jax.lax.acos
jax.numpy.arccos
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Acosh ✅ ✅ ✅ jax/lax/acosh
jax/numpy/acosh
jax.lax.acosh
jax.numpy.arccosh
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Add ✅ ✅ ✅ _normalization_utils
equinox/eqx/nn/recurrent
equinox/eqx/nn/rms_norm
equinox/eqx/nn/rotary_positional_embedding
flax/linen/dot_product_attention
flax/linen/einsum
... +70
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
AffineGrid ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
And ✅ ✅ ✅ _normalization_utils
flax/nnx/combine_masks
flax/nnx/group_norm
jax/lax/and
jax/lax/betainc
jax/lax/clz
... +11
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ArgMax ✅ ✅ ✅ jax/lax/_arg_utils
jax/lax/argmax
jax/lax/lu
jax/numpy/argmax
jax/numpy/composite_metadata_batch3
jax/random/categorical
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ArgMin ✅ ✅ ✅ jax/lax/_arg_utils
jax/lax/argmin
jax/numpy/argmin
jax/numpy/composite_metadata_batch3
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Asin ✅ ✅ ✅ jax/lax/asin
jax/numpy/asin
jax.lax.asin
jax.numpy.arcsin
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Asinh ✅ ✅ ✅ jax/lax/asinh
jax/numpy/asinh
jax.lax.asinh
jax.numpy.arcsinh
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Atan ✅ ✅ ✅ jax/lax/atan
jax/lax/atan2
jax/numpy/atan
jax/numpy/atan2
jax.lax.atan
jax.numpy.arctan
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Atanh ✅ ✅ ✅ jax/lax/atanh
jax/numpy/atanh
jax.lax.atanh
jax.numpy.arctanh
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Attention ✅ ✅ ✅ equinox/eqx/nn/multihead_attention
flax/nnx/dot_product_attention
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
AveragePool ✅ ✅ ✅ equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/avg_pool
flax/linen/avg_pool
flax/linen/pool
flax/nnx/avg_pool
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BatchNormalization ✅ ✅ ✅ equinox/eqx/nn/batch_norm
flax/linen/batch_norm
flax/nnx/batch_norm
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Bernoulli ✅ ✅ ✅ jax/random/bernoulli Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BitCast ✅ ✅ ✅ jax/lax/bitcast_convert_type Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BitShift ✅ ✅ ✅ jax/lax/clz
jax/lax/population_count
jax/lax/shift_left
jax/lax/shift_right_arithmetic
jax/lax/shift_right_logical
jax/numpy/bitwise_left_shift
... +3
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BitwiseAnd ✅ ✅ ✅ jax/lax/and
jax/lax/clz
jax/lax/population_count
jax/numpy/bitwise_and
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BitwiseNot ✅ ✅ ✅ jax/lax/bitwise_not
jax/numpy/bitwise_not
jax/numpy/invert
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BitwiseOr ✅ ✅ ✅ jax/lax/clz
jax/lax/or
jax/lax/shift_right_arithmetic
jax/numpy/bitwise_or
jax/numpy/right_shift
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BitwiseXor ✅ ✅ ✅ jax/lax/xor
jax/numpy/bitwise_xor
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
BlackmanWindow ✅ ✅ ✅ jax/numpy/windows Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Cast ✅ ✅ ✅ _complex_utils
_normalization_utils
_precision_utils
equinox/eqx/nn/embedding
flax/linen/dot_product_attention
flax/nnx/combine_masks
... +84
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
CastLike ✅ ✅ ✅ _utils
flax/linen/conv
flax/linen/dense_general
flax/nnx/conv
flax/nnx/group_norm
flax/nnx/linear
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
CausalConvWithState ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Evaluate demand and either implement plugin support or document non-goal.
Ceil ✅ ✅ ✅ jax/lax/ceil
jax/numpy/ceil
jax.lax.ceil
jax.numpy.ceil
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Celu ✅ ✅ ✅ jax/nn/celu Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
CenterCropPad ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
Clip ✅ ✅ ✅ flax/nnx/hard_tanh
jax/lax/convert_element_type
jax/numpy/clip
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Col2Im ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
Compress ✅ ✅ ✅ jax/numpy/compress Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Concat ✅ ✅ ✅ _attention_utils
_axis0_utils
_complex_utils
equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/linear
equinox/eqx/nn/multihead_attention
... +42
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ConcatFromSequence ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
Constant ✅ ✅ ✅ flax/nnx/dropout
jax/lax/empty
jax/numpy/windows
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ConstantOfShape ✅ ✅ ✅ jax/lax/dot_general
jax/numpy/eye
jax/numpy/full
jax/numpy/ones
jax/numpy/zeros
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Conv ✅ ✅ ✅ equinox/eqx/nn/conv
flax/linen/conv
flax/linen/conv_local
flax/linen/recurrent
flax/nnx/conv
jax/lax/conv
... +2
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ConvInteger ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Decide quantization scope, then add lowerings/tests or mark as not planned.
ConvTranspose ✅ ✅ ✅ flax/linen/conv_transpose
jax/lax/conv
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Cos ✅ ✅ ✅ jax/lax/cos
jax/lax/digamma
jax/numpy/cos
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Cosh ✅ ✅ ✅ jax/lax/cosh
jax/numpy/cosh
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
CumProd ✅ ✅ ✅ jax/lax/cumprod
jax/numpy/cumprod
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
CumSum ✅ ✅ ✅ jax/lax/cumlogsumexp
jax/lax/cumsum
jax/numpy/composite_metadata_batch3
jax/numpy/cumsum
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
DFT ✅ ✅ ✅ jax/lax/fft
jax/numpy/fft
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
DeformConv ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
DepthToSpace ✅ ✅ ✅ dm_pix/depth_to_space Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
DequantizeLinear ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Decide quantization scope, then add lowerings/tests or mark as not planned.
Det ✅ ✅ ✅ jax/numpy/linalg_det Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Div ✅ ✅ ✅ _normalization_utils
equinox/eqx/nn/rms_norm
equinox/eqx/nn/spectral_norm
equinox/eqx/nn/weight_norm
flax/linen/dot_product_attention
flax/linen/spectral_norm
... +55
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Dropout ✅ ✅ ✅ equinox/eqx/nn/dropout
flax/linen/dropout
flax/nnx/dot_product_attention
flax/nnx/dropout
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
DynamicQuantizeLinear ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Decide quantization scope, then add lowerings/tests or mark as not planned.
Einsum ✅ ✅ ✅ flax/linen/einsum
flax/nnx/einsum
jax/lax/dot_general
jax/nn/scaled_dot_general
jax/numpy/einsum
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Elu ✅ ✅ ✅ flax/nnx/elu
jax/nn/elu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Equal ✅ ✅ ✅ _normalization_utils
flax/nnx/group_norm
jax/lax/_reduce_utils
jax/lax/atan2
jax/lax/betainc
jax/lax/clz
... +27
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Erf ✅ ✅ ✅ flax/nnx/gelu
jax/lax/erf
jax/lax/erf_inv
jax/lax/erfc
jax/nn/gelu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Exp ✅ ✅ ✅ jax/lax/bessel_i0e
jax/lax/bessel_i1e
jax/lax/betainc
jax/lax/cosh
jax/lax/cumlogsumexp
jax/lax/erf_inv
... +12
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Expand ✅ ✅ ✅ _axis0_utils
flax/linen/group_norm
jax/lax/broadcast_in_dim
jax/lax/dynamic_update_slice
jax/lax/gather
jax/lax/iota
... +6
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
EyeLike ✅ ✅ ✅ jax/numpy/eye Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Flatten ✅ ✅ ✅ jax/numpy/reshape Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Floor ✅ ✅ ✅ jax/lax/floor
jax/lax/nextafter
jax/lax/polygamma
jax/lax/reduce_precision
jax/numpy/floor
jax/numpy/floor_divide
... +3
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GRU ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
Gather ✅ ✅ ✅ _complex_utils
equinox/eqx/nn/embedding
equinox/eqx/nn/multihead_attention
flax/linen/embed
flax/nnx/conv
flax/nnx/embed
... +42
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GatherElements ✅ ✅ ✅ flax/nnx/flip_sequences
jax/lax/_sort_utils
jax/lax/sort
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GatherND ✅ ✅ ✅ jax/lax/gather Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Gelu ✅ ✅ ✅ flax/nnx/gelu
jax/nn/gelu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Gemm ✅ ✅ ✅ equinox/eqx/nn/linear
equinox/eqx/nn/multihead_attention
equinox/eqx/nn/recurrent
flax/linen/conv_local
flax/linen/dense
flax/linen/dense_general
... +8
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GlobalAveragePool ✅ ✅ ✅ flax/linen/avg_pool
flax/nnx/avg_pool
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GlobalLpPool ✅ ✅ ✅ jax/numpy/linalg_norm Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GlobalMaxPool ✅ ✅ ✅ flax/linen/max_pool
flax/nnx/max_pool
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Greater ✅ ✅ ✅ jax/lax/atan2
jax/lax/betainc
jax/lax/digamma
jax/lax/gather
jax/lax/gt
jax/lax/igamma
... +9
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GreaterOrEqual ✅ ✅ ✅ jax/lax/ge
jax/numpy/digitize
jax/numpy/greater_equal
jax/numpy/histogram
jax/numpy/histogram2d
jax/numpy/histogramdd
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GridSample ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
GroupNormalization ✅ ✅ ✅ flax/nnx/group_norm Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
GroupedMatMul ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Evaluate demand and either implement plugin support or document non-goal.
HammingWindow ✅ ✅ ✅ jax/numpy/windows Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
HannWindow ✅ ✅ ✅ jax/numpy/windows Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
HardSigmoid ✅ ✅ ✅ jax/nn/hard_sigmoid Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
HardSwish ✅ ✅ ✅ jax/nn/hard_swish Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Hardmax ✅ ✅ ✅ jax/nn/hardmax Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Identity ✅ ✅ ✅ equinox/eqx/nn/identity
equinox/eqx/nn/rms_norm
equinox/eqx/nn/sequential
equinox/internal/control_primitives
jax/core/name
jax/lax/_control_flow_utils
... +54
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
If ✅ ✅ ✅ jax/lax/cond Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ImageDecoder ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
InstanceNormalization ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Evaluate demand and either implement plugin support or document non-goal.
IsInf ✅ ✅ ✅ jax/lax/is_finite
jax/lax/nextafter
jax/numpy/frexp
jax/numpy/isfinite
jax/numpy/spacing
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
IsNaN ✅ ✅ ✅ jax/lax/_sort_utils
jax/lax/convert_element_type
jax/lax/is_finite
jax/lax/nextafter
jax/numpy/composite_metadata_batch2
jax/numpy/composite_metadata_batch3
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LRN ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
LSTM ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
LayerNormalization ✅ ✅ ✅ _normalization_utils
equinox/eqx/nn/layer_norm
flax/linen/layer_norm
flax/nnx/layer_norm
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LeakyRelu ✅ ✅ ✅ flax/nnx/leaky_relu
jax/nn/leaky_relu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Less ✅ ✅ ✅ flax/linen/dot_product_attention
jax/lax/atan2
jax/lax/bessel_i0e
jax/lax/bessel_i1e
jax/lax/betainc
jax/lax/digamma
... +29
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LessOrEqual ✅ ✅ ✅ jax/lax/eigh
jax/lax/le
jax/numpy/composite_metadata_batch2
jax/numpy/digitize
jax/numpy/histogram
jax/numpy/histogram2d
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LinearAttention ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Evaluate demand and either implement plugin support or document non-goal.
Log ✅ ✅ ✅ jax/lax/cumlogsumexp
jax/lax/digamma
jax/lax/erf_inv
jax/lax/expm1
jax/lax/lgamma
jax/lax/log
... +13
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LogSoftmax ✅ ✅ ✅ flax/nnx/log_softmax
jax/nn/log_softmax
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Loop ✅ ✅ ✅ flax/linen/recurrent
jax/lax/_control_flow_utils
jax/lax/fori_loop
jax/lax/scan
jax/lax/while_loop
jax/numpy/composite_metadata_batch6
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LpNormalization ✅ ✅ ✅ jax/lax/div Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
LpPool ✅ ✅ ✅ jax/lax/reduce_window_sum Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
MatMul ✅ ✅ ✅ equinox/eqx/nn/multihead_attention
equinox/eqx/nn/spectral_norm
flax/linen/dot_product_attention
flax/linen/multi_head_attention
flax/linen/multi_head_dot_product_attention
flax/linen/self_attention
... +20
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
MatMulInteger ✅ ✅ ✅ jax/lax/dot_general Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Max ✅ ✅ ✅ _normalization_utils
flax/nnx/group_norm
jax/lax/clamp
jax/lax/eig
jax/lax/gather
jax/lax/igamma
... +9
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
MaxPool ✅ ✅ ✅ equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/max_pool
equinox/eqx/nn/pool
flax/linen/max_pool
flax/linen/min_pool
flax/linen/pool
... +6
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
MaxRoiPool ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
MaxUnpool ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
Mean ✅ ✅ ✅ jax/lax/div Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
MeanVarianceNormalization ✅ ✅ ✅ jax/nn/standardize Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
MelWeightMatrix ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
Min ✅ ✅ ✅ jax/lax/clamp
jax/lax/gather
jax/lax/igamma
jax/lax/min
jax/lax/shift_right_arithmetic
jax/nn/hard_tanh
... +6
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Mish ✅ ✅ ✅ jax/nn/mish Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Mod ✅ ✅ ✅ jax/lax/_reduce_utils
jax/lax/reduce_xor
jax/lax/rem
jax/numpy/composite_metadata_batch2
jax/numpy/composite_metadata_batch3
jax/numpy/fmod
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Mul ✅ ✅ ✅ _complex_utils
_normalization_utils
equinox/eqx/nn/multihead_attention
equinox/eqx/nn/recurrent
equinox/eqx/nn/rms_norm
equinox/eqx/nn/rotary_positional_embedding
... +91
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Multinomial ✅ ✅ ✅ jax/random/categorical Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Neg ✅ ✅ ✅ _complex_utils
equinox/eqx/nn/rotary_positional_embedding
flax/linen/min_pool
flax/linen/pool
flax/nnx/log_sigmoid
jax/lax/_cum_extrema
... +17
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
NegativeLogLikelihoodLoss ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
NonMaxSuppression ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
NonZero ✅ ✅ ✅ jax/lax/scatter
jax/lax/scatter_utils
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Not ✅ ✅ ✅ equinox/eqx/nn/dropout
flax/nnx/dot_product_attention
flax/nnx/dropout
jax/lax/betainc
jax/lax/bitwise_not
jax/lax/igamma
... +9
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
OneHot ✅ ✅ ✅ jax/nn/one_hot Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Optional ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
OptionalGetElement ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
OptionalHasElement ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
Or ✅ ✅ ✅ jax/lax/atan2
jax/lax/betainc
jax/lax/expm1
jax/lax/gather
jax/lax/is_finite
jax/lax/log1p
... +14
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
PRelu ✅ ✅ ✅ equinox/eqx/nn/prelu
flax/linen/prelu
flax/nnx/prelu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Pad ✅ ✅ ✅ _axis0_utils
jax/lax/conv
jax/lax/dot_general
jax/lax/pad
jax/lax/reduce_window_sum
jax/numpy/composite_metadata_batch4
... +1
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Pow ✅ ✅ ✅ jax/lax/betainc
jax/lax/cbrt
jax/lax/exp2
jax/lax/igamma
jax/lax/integer_pow
jax/lax/nextafter
... +10
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
QLinearConv ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Decide quantization scope, then add lowerings/tests or mark as not planned.
QLinearMatMul ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Decide quantization scope, then add lowerings/tests or mark as not planned.
QuantizeLinear ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Decide quantization scope, then add lowerings/tests or mark as not planned.
RMSNormalization ✅ ✅ ✅ equinox/eqx/nn/rms_norm
flax/nnx/rms_norm
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
RNN ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
RandomNormal ✅ ✅ ✅ jax/random/normal Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
RandomNormalLike ✅ ✅ ✅ jax/random/normal Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
RandomUniform ✅ ✅ ✅ jax/lax/rng_bit_generator
jax/lax/rng_uniform
jax/random/random_bits
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
RandomUniformLike ✅ ✅ ✅ jax/random/categorical
jax/random/random_bits
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Range ✅ ✅ ✅ jax/lax/dynamic_update_slice
jax/lax/gather
jax/lax/iota
jax/lax/rev
jax/numpy/arange
jax/numpy/composite_metadata
... +3
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Reciprocal ✅ ✅ ✅ jax/lax/integer_pow jax.numpy.reciprocal
jax.lax.integer_pow(x, -1)
1.0 / x
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceL1 ✅ ✅ ✅ jax/lax/reduce_sum
jax/numpy/sum
jax.numpy.linalg.norm(ord=1)
jax.numpy.sum(jax.numpy.abs(x), axis=...)
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceL2 ✅ ✅ ✅ jax/lax/sqrt
jax/numpy/sqrt
jax.numpy.linalg.norm(ord=2)
jax.numpy.sqrt(jax.numpy.sum(x * x, axis=...))
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceLogSum ✅ ✅ ✅ jax/lax/log
jax/numpy/log
jax.numpy.log(jax.numpy.sum(x, axis=...)) Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceLogSumExp ✅ ✅ ✅ jax/lax/log
jax/nn/logmeanexp
jax/nn/logsumexp
jax/numpy/log
jax.nn.logsumexp
jax.scipy.special.logsumexp
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceMax ✅ ✅ ✅ _normalization_utils
equinox/eqx/nn/adaptive_pool
equinox/internal/control_primitives
flax/nnx/group_norm
jax/lax/reduce
jax/lax/reduce_max
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceMean ✅ ✅ ✅ _normalization_utils
equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/rms_norm
flax/nnx/group_norm
flax/nnx/rms_norm
jax/numpy/composite_metadata
... +3
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceMin ✅ ✅ ✅ _normalization_utils
equinox/internal/control_primitives
flax/nnx/group_norm
jax/lax/_reduce_utils
jax/lax/reduce
jax/lax/reduce_and
... +6
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceProd ✅ ✅ ✅ flax/nnx/conv
flax/nnx/group_norm
jax/lax/_cum_extrema
jax/lax/dot_general
jax/lax/dynamic_update_slice
jax/lax/reduce
... +5
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceSum ✅ ✅ ✅ equinox/internal/control_primitives
flax/linen/dot_product_attention
flax/linen/group_norm
jax/lax/_reduce_utils
jax/lax/reduce
jax/lax/reduce_or
... +16
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReduceSumSquare ✅ ✅ ✅ flax/linen/group_norm
jax/lax/reduce_sum
jax/numpy/sum
jax.numpy.sum(jax.numpy.square(x), axis=...) Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
RegexFullMatch ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
Relu ✅ ✅ ✅ equinox/eqx/nn/lambda
equinox/eqx/nn/sequential
flax/nnx/relu
jax/nn/relu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Reshape ✅ ✅ ✅ _attention_utils
equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/linear
equinox/eqx/nn/multihead_attention
flax/linen/conv
flax/linen/conv_local
... +42
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Resize ✅ ✅ ✅ jax/image/resize Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ReverseSequence ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
RoiAlign ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Vision-specific op; add only when a target model needs the native ONNX op.
RotaryEmbedding ✅ ✅ ✅ equinox/eqx/nn/rotary_positional_embedding Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Round ✅ ✅ ✅ jax/lax/polygamma
jax/lax/reduce_precision
jax/lax/round
jax.lax.round
jax.numpy.round
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
STFT ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
Scan ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
Scatter ✅ ✅ ✅ jax/lax/scatter
jax/lax/scatter_utils
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ScatterElements ✅ ✅ ✅ jax/lax/lu
jax/lax/lu_pivots_to_permutation
jax/lax/scatter
jax/lax/scatter_utils
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
ScatterND ✅ ✅ ✅ jax/lax/cholesky
jax/lax/cholesky_update
jax/lax/dynamic_update_slice
jax/lax/hessenberg
jax/lax/householder_product
jax/lax/lu
... +11
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Selu ✅ ✅ ✅ jax/nn/selu Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
SequenceAt ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
SequenceConstruct ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
SequenceEmpty ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
SequenceErase ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
SequenceInsert ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
SequenceLength ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
SequenceMap ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
Shape ✅ ✅ ✅ _axis0_utils
equinox/eqx/nn/linear
equinox/eqx/nn/multihead_attention
flax/linen/dense
flax/linen/dense_general
flax/nnx/conv
... +15
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Shrink ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Evaluate demand and either implement plugin support or document non-goal.
Sigmoid ✅ ✅ ✅ equinox/eqx/nn/lambda
equinox/eqx/nn/recurrent
equinox/eqx/nn/sequential
flax/linen/recurrent
flax/nnx/glu
flax/nnx/sigmoid
... +5
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Sign ✅ ✅ ✅ jax/lax/bessel_i1e
jax/lax/cbrt
jax/lax/eig
jax/lax/erf_inv
jax/lax/nextafter
jax/lax/reduce_precision
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Sin ✅ ✅ ✅ jax/lax/cos
jax/lax/digamma
jax/lax/lgamma
jax/lax/sin
jax/numpy/composite_metadata_batch2
jax/numpy/cos
... +2
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Sinh ✅ ✅ ✅ jax/lax/sinh Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Size ✅ ✅ ✅ jax/numpy/size Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Slice ✅ ✅ ✅ _axis0_utils
equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/linear
equinox/eqx/nn/multihead_attention
flax/linen/conv_transpose
flax/linen/dense
... +21
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Softmax ✅ ✅ ✅ equinox/eqx/nn/multihead_attention
flax/linen/dot_product_attention
flax/linen/multi_head_attention
flax/linen/multi_head_dot_product_attention
flax/linen/self_attention
flax/nnx/dot_product_attention
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
SoftmaxCrossEntropyLoss ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
Softplus ✅ ✅ ✅ flax/nnx/log_sigmoid
flax/nnx/softplus
jax/nn/softplus
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Softsign ✅ ✅ ✅ jax/nn/softsign Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
SpaceToDepth ✅ ✅ ✅ dm_pix/space_to_depth Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Split ✅ ✅ ✅ equinox/eqx/nn/rotary_positional_embedding
flax/nnx/glu
jax/lax/split
jax/nn/glu
jax/numpy/split
jax/numpy/unstack
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
SplitToSequence ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
Sqrt ✅ ✅ ✅ _normalization_utils
equinox/eqx/nn/rms_norm
flax/nnx/group_norm
flax/nnx/rms_norm
jax/lax/bessel_i0e
jax/lax/bessel_i1e
... +15
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Squeeze ✅ ✅ ✅ dm_pix/depth_to_space
dm_pix/space_to_depth
equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/avg_pool
equinox/eqx/nn/conv
equinox/eqx/nn/max_pool
... +16
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
StringConcat ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
StringNormalizer ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
StringSplit ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. If in scope, add container plugins; else mark explicitly out-of-scope.
Sub ✅ ✅ ✅ _normalization_utils
flax/nnx/group_norm
jax/lax/atan2
jax/lax/betainc
jax/lax/cholesky
jax/lax/cholesky_update
... +53
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Sum ✅ ✅ ✅ jax/lax/add_any Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
SwiGLU ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Evaluate demand and either implement plugin support or document non-goal.
Swish ✅ ✅ ✅ flax/nnx/silu
jax/nn/silu
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Tan ✅ ✅ ✅ jax/lax/tan
jax/numpy/tan
jax.lax.tan
jax.numpy.tan
Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Tanh ✅ ✅ ✅ equinox/eqx/nn/lambda
equinox/eqx/nn/recurrent
equinox/eqx/nn/sequential
flax/linen/recurrent
flax/nnx/gelu
flax/nnx/tanh
... +4
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
TensorScatter ✅ ✅ ✅ jax/lax/dynamic_update_slice Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
TfIdfVectorizer ➖ ➖ ➖ ➖ No clear direct primitive; likely composite JAX expression. Add only when demanded by target models; document priority.
ThresholdedRelu ✅ ✅ ✅ jax/nn/thresholded_relu Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Tile ✅ ✅ ✅ _attention_utils
jax/nn/scaled_matmul
jax/numpy/tile
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
TopK ✅ ✅ ✅ jax/lax/approx_top_k
jax/lax/sort
jax/lax/top_k
jax/numpy/composite_metadata_batch2
jax/numpy/sort
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Transpose ✅ ✅ ✅ dm_pix/depth_to_space
dm_pix/space_to_depth
equinox/eqx/nn/multihead_attention
flax/linen/avg_pool
flax/linen/conv
flax/linen/conv_local
... +34
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Trilu ✅ ✅ ✅ jax/numpy/trilu Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Unique ✅ ✅ ✅ jax/numpy/unique Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Unsqueeze ✅ ✅ ✅ _attention_utils
_complex_utils
dm_pix/depth_to_space
dm_pix/space_to_depth
equinox/eqx/nn/adaptive_pool
equinox/eqx/nn/avg_pool
... +40
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Upsample ✅ ✅ ✅ jax/image/resize Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Where ✅ ✅ ✅ _normalization_utils
flax/linen/dot_product_attention
flax/nnx/dot_product_attention
flax/nnx/group_norm
jax/lax/_sort_utils
jax/lax/atan2
... +45
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.
Xor ✅ ✅ ✅ jax/lax/xor
jax/numpy/bitwise_xor
jax/numpy/floor_divide
Covered via existing plugin primitives. Keep covered; broaden tests (dynamic shapes/dtypes/opset) and sync metadata.

Plugin Names Not In Index

This section lists operator-like names found in plugins that do not appear in the ONNX index page.

Name Found in Plugins Metadata Lowering Plugin Modules Potential JAX Ops Next Action
Cholesky ✅ ➖ jax/numpy/composite_metadata_batch4 Internal alias/helper; map manually if needed. Metadata alias only; replace with official ONNX operator name or remove.
Complex ➖ ✅ _complex_utils jax.lax.complex
jax.numpy.asarray(..., dtype=complex64/128)
Internal helper op composition; keep as helper, not as ONNX coverage target.
Eig ✅ ➖ jax/numpy/composite_metadata_batch5 Internal alias/helper; map manually if needed. Metadata alias only; replace with official ONNX operator name or remove.
Eigh ✅ ➖ jax/numpy/composite_metadata_batch5 Internal alias/helper; map manually if needed. Metadata alias only; replace with official ONNX operator name or remove.
QR ✅ ➖ jax/numpy/composite_metadata_batch4 Internal alias/helper; map manually if needed. Metadata alias only; replace with official ONNX operator name or remove.
SVD ✅ ➖ jax/numpy/composite_metadata_batch4
jax/numpy/composite_metadata_batch5
Internal alias/helper; map manually if needed. Metadata alias only; replace with official ONNX operator name or remove.