Feature request: opset 23+ support

[justinchuby]

Opset 23 introduced Attention and RotaryEmbedding operators, which can simplify some of the subgraphs generated from corresponding equinox ops.

[enpasos]

Agreed — opset 23’s Attention/RotaryEmbedding would be a nice simplification. If ORT 1.23.1 (CPU) already has kernels for both, we can emit the fused ops for opset ≥23 and add tests; we’ll keep the pre-23 decomposition as a fallback to support older runtimes.

[justinchuby]

Yes it does have kernels for both.

[enpasos]

I gave it a try on branch opset23. The exporter can emit the fused-ops graph (tested on a DINOv3 VisionTransformer), but I hit what looks like an onnxruntime issue in shape/type inference when nested custom FunctionProtos are present. Any idea?

---

Problem Summary

Emitting the DINOv3 VisionTransformer at opset 23 produces nested custom FunctionProtos (VisionTransformer → Block → LayerScale/MultiHeadAttention).
With those nested function domains intact, onnxruntime 1.23.2 fails during shape inference (and even earlier misinterprets RotaryEmbedding inputs), repeatedly reporting:

(op_type:RotaryEmbedding, node name: node_RotaryEmbedding_24): [ShapeInferenceError] Input shape is 3D, num_heads attribute must be provided
(op_type:Reshape …): [TypeInferenceError] Input 0 expected to have type but instead is null
(op_type:LayerScale …): [TypeInferenceError] Input 0 expected to have type but instead is null

Exporting the same DINOv3 graph for opset 21 does not hit the issue; shape inference and an onnxruntime.InferenceSession both succeed.
Conclusion: the regression appears confined to the new fused-ops path at opset 23.

---

Proof-of-Failure / Repro

From the repo root, run:

poetry run python jax2onnx/sandbox/ort_nested_function_failure.py

The script:

• Prints the onnxruntime version (1.23.2).
• Loads docs/onnx/examples/eqx_dino/eqx_dinov3_vit_S16_opset23.onnx.
• Runs onnx.shape_inference.
• Surfaces the same “Input … expected to have type but instead is null” error chain caused by nested custom FunctionProtos.

[justinchuby]

Thanks, I will create an issue on onnx runtime. This may be due to its function support because onnx functions are not best supported by ORT