Skip to main content

Module 9 — Unsupported operators and workarounds

Up to this point every export ran to completion. On real production models this stops being true around the third or fourth exotic architecture: a model uses a function that PyTorch or TensorFlow implements natively, that ONNX either does not define at all or defines only in a newer opset than the exporter targets, and the export raises. This module is about what to do at that moment. Nine times out of ten the fix is not "wait for a new ONNX release" but a small, surgical change on the model side.

The text encoder in our running example is deliberately built to hit this problem. Its attention mask is computed with a Python trick that traces poorly, and one of its normalisation variants uses an operator absent from opset 13. Both are common issues, both have known workarounds, and both are described here.

Reading the error​

The first useful thing to do when an export raises is to read the message carefully. PyTorch, TensorFlow and tf2onnx all print the offending operator's name and its location in the model:

RuntimeError: Exporting the operator 'aten::scaled_dot_product_attention'
    to ONNX opset version 11 is not supported. Please feel free to request
    support or submit a pull request on PyTorch GitHub:
    https://github.com/pytorch/pytorch/issues.

Three pieces of information matter here. aten::scaled_dot_product_attention identifies the operator. opset version 11 identifies what was requested. The remaining text is boilerplate. The same operator is supported on opset 14 and above; the fix is one line — change opset_version=11 to opset_version=17.

Some messages point at a specific location in the model instead of a generic operator:

Unable to lower expression at graph.py:214, node 'my_attention.forward'

Line 214 of the traceback is where to look. The pattern is almost always a Python operation on tensor values — if x.item() > 0:, for i in range(x.shape[0]):, x.tolist() — that the tracer cannot serialise.

Four workarounds in order of preference​

Not all workarounds are equal. Prefer them in this order, from safest to most invasive:

1. Change the opset. If the operator is supported on a newer opset, raise opset_version in torch.onnx.export. Confirm the ONNX Runtime version at deployment can execute the new opset (Runtime 1.17+ supports opset 20, older runtimes stop earlier). Rerun module 4's parity check.

2. Use the dynamo exporter. torch.onnx.export(..., dynamo=True) handles more Python control flow than the legacy tracer and often exports transformer attention successfully where the tracer refuses.

3. Rewrite the offending module. If the failure is on Python control flow — an if on a tensor value, a Python for-loop — rewrite the module using tensor operations. torch.where(condition, a, b) replaces if; broadcasting replaces per-element loops; masked_fill replaces conditional zeroing. This is where most exporter bugs get resolved.

4. Register a custom operator. Reserved for genuinely novel operations, for example a specialised loss or a hardware-specific kernel. The mechanism is documented but requires C++ knowledge and a runtime rebuild for exotic providers; use only when the first three fail.

The text encoder: attention mask that traces poorly​

The naive text encoder computes its attention mask on the fly:

class TextEncoder(torch.nn.Module):
    def forward(self, input_ids):
        # PAD token is 0
        mask = (input_ids != 0)
        seq_len = mask.sum(dim=1).max().item()      # traces poorly
        embedded = self.embed(input_ids)[:, :seq_len]
        ...

The .item() call converts a tensor to a Python scalar and destroys traceability — the tracer can only record what stays inside the tensor world. The exporter either refuses outright or emits a graph where seq_len is frozen to the value of the example input. Either way the export fails.

The clean rewrite keeps everything in tensor land:

class TextEncoder(torch.nn.Module):
    def forward(self, input_ids, attention_mask):
        # attention_mask is now an explicit input, passed by the caller
        embedded = self.embed(input_ids)
        return self.transformer(embedded, src_key_padding_mask=~attention_mask.bool())

Two changes: attention_mask becomes an explicit input, and the runtime slicing disappears. The export produces a graph with two dynamic axes on both inputs, and inference at any sequence length works. The caller now passes both input_ids and attention_mask, matching the standard transformer API of Hugging Face and other frameworks — a small alignment cost that is worth paying anyway.

The custom LayerNormalization case​

Older opsets do not define LayerNormalization as a single operator. Before opset 17, the exporter would decompose it into ReduceMean, ReduceMean, Sub, Pow, Sqrt, Div, Mul, Add — eight primitive operators for one logical operation. The resulting graph runs, is numerically correct, and is 2 to 3x slower than the fused version, because runtimes cannot easily recognise the pattern and refuse the specialised kernel.

The fix here is trivial: opset_version=17 produces a single LayerNormalization node, which every runtime executes as a fused kernel. Every checkbox — export, parity, optimization — flips green.

The general lesson: newer opsets are almost never a downgrade. Unless a specific target runtime constrains you, prefer the newest supported opset.

Registering a custom operator​

The case where none of the above work is a genuinely novel operator — a research paper's brand-new activation, a bespoke tokenisation step, a hardware-accelerated primitive. ONNX supports custom operators through a two-step registration.

On the PyTorch side, register a symbolic function that tells the exporter how to translate the operator:

import torch
from torch.onnx import register_custom_op_symbolic

def _my_op(g, x, alpha):
    return g.op("mynamespace::MyOp", x, alpha_f=alpha)

register_custom_op_symbolic("mynamespace::my_op", _my_op, 17)

On the runtime side, register a matching kernel implementation. In Python via onnxruntime-extensions, or in C++ against the ONNX Runtime plugin API. The full path is documented in the ONNX Runtime custom-op guide.

The trade-off is real. A custom operator escapes the promise "any runtime can load this model". Whoever consumes the file must also have the custom-op library installed at exactly the compatible version. This voids most of the reason to use ONNX in the first place, so exhaust workarounds 1 to 3 before reaching for a custom op.

Two-thirds of unsupported-operator errors disappear with opset_version=17

Whenever an export raises, before rewriting anything, bump the opset to 17 (or 18/19 if your runtime supports it) and try again. LayerNormalization, GroupNormalization, ScaledDotProductAttention, RotaryEmbedding and a dozen others became native operators over the last three opset generations. The default in older tutorials is opset 11, which is why so many "unsupported operator" bug reports become one-line fixes.

From TensorFlow, the same problem in a different accent​

tf2onnx fails with a similar message shape:

tf2onnx.tfonnx - ERROR - Unsupported ops: Counter({'HashTable': 1, 'LookupTableFindV2': 1})

The two operators named here — HashTable and LookupTableFindV2 — belong to text preprocessing: a vocabulary lookup baked into the Keras model. The workaround is to cut the preprocessing out of the model. The vocabulary lookup lives in Python, the tokenised integer tensor is what enters the ONNX graph. This mirrors the PyTorch text encoder rewrite: preprocessing outside, pure tensor operations inside.

In summary​

  • The error message names the operator and the opset; both are actionable pieces of information, not boilerplate.
  • Prefer the workarounds in this order: raise the opset, try the dynamo exporter, rewrite the module to remove Python-side control flow, register a custom operator as a last resort.
  • The text encoder illustrates the classic transformer trap: an attention mask computed with .item() traces poorly; pass the mask as an explicit input and everything downstream falls into place.
  • Preprocessing that uses vocabulary lookups or string tables does not belong inside the ONNX graph — cut it out and keep the model to pure numeric tensor operations.

Next module: putting all of this behind a real serving endpoint.