Skip to content

onnx: add LpNormalization, MeanVarianceNormalization, GroupNormalization, RotaryEmbedding op handlers#2283

Merged
kali merged 1 commit into
sonos:mainfrom
czoli1976:feature/onnx-import-bindings
May 26, 2026
Merged

onnx: add LpNormalization, MeanVarianceNormalization, GroupNormalization, RotaryEmbedding op handlers#2283
kali merged 1 commit into
sonos:mainfrom
czoli1976:feature/onnx-import-bindings

Conversation

@czoli1976
Copy link
Copy Markdown
Contributor

Summary

Adds ONNX importer handlers for four inference-relevant operators that tract could already compute as decomposed subgraphs but had no single-node binding for — continuing the recent Attention/RMSNorm/Gelu/Swish/Mish additions:

  • LpNormalization (opset 1) — Reduce<L1|L2> over axis, then divide
  • MeanVarianceNormalization (opset 13) — standard function decomposition (epsilon outside the sqrt)
  • GroupNormalization (opset 18 & 21) — channel-split normalization; opset-aware affine ([num_groups] pre-21, [C] from 21) and honors stash_type (f32 reduction for f16/bf16 inputs)
  • RotaryEmbedding (opset 23) — full reference algorithm: 3D/4D input, optional position_ids, partial rotary_embedding_dim, NeoX and GPT-J (interleaved) layouts

All lower to existing tract primitives — importer coverage, no new kernels.

Tests

  • Corresponding ONNX backend node tests enabled in test-rt/suite-onnx/node.txt, green on default + unoptimized runtimes (GroupNorm at both the opset-18 and opset-21 affine paths; all 8 RoPE variants). Full node suite: 4428 passed / 0 failed.
  • Spot-checked end-to-end against ONNX ReferenceEvaluator on small multi-op graphs, including dynamic (symbolic) batch/sequence shapes; RoPE bit-exact.

Note

tract resolves ops by name regardless of domain; this PR targets the standardized ai.onnx operators. The com.microsoft RoPE variant has a different input order and is intentionally out of scope here.

…ion, RotaryEmbedding op handlers

Import handlers for four standardized ai.onnx operators, lowering to
existing tract primitives:
- LpNormalization (opset 1), MeanVarianceNormalization (opset 13)
- GroupNormalization (opset 18 & 21; opset-aware affine, f32 stash_type)
- RotaryEmbedding (opset 23; 3D/4D input, position_ids, partial + interleaved)

Corresponding ONNX backend node tests enabled in suite-onnx/node.txt.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@kali kali merged commit 4fa823e into sonos:main May 26, 2026
55 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants