latest (dev)
Copy
Latest development documentation · Updated 2026-10-15
Source code for tensorplay.onnx._composite_ops
"""ONNX lowering for TensorPlay graph nodes.
Every captured ``call_function`` / ``call_method`` node is translated by a
handler registered here. Handlers receive an :class:`OpContext`, read their
arguments by *name* (so positional and keyword capture behave identically) and
emit ONNX nodes through the :class:`GraphBuilder`.
Ops that ONNX expresses directly become a single node; the rest are lowered to
an equivalent ONNX subgraph (``linear`` -> ``Gemm``/``MatMul+Add``, ``silu`` ->
``Mul(x, Sigmoid(x))``, ``group_norm`` -> explicit moments, ...).
"""
from __future__ import annotations
import math
from typing import Any, Callable, Sequence
import numpy as np
from onnx import helper, numpy_helper
from ._type_mapping import _dtype_to_onnx, _np_dtype_to_onnx, _to_numpy
__all__ = [
"GraphBuilder",
"OpContext",
"Value",
"lookup_function_handler",
"lookup_method_handler",
]
_INT64_MAX = np.iinfo(np.int64).max
_INT64_MIN = np.iinfo(np.int64).min
from .errors import UnsupportedOperatorError # noqa: F401
# ---------------------------------------------------------------------------
# Values and graph building
# ---------------------------------------------------------------------------
[docs]
class Value:
"""A tensor flowing through the ONNX graph under construction."""
__slots__ = ("name", "shape", "dtype")
def __init__(
self,
name: str,
shape: tuple | None = None,
dtype: np.dtype | None = None,
) -> None:
self.name = name
self.shape = shape
self.dtype = np.dtype(dtype) if dtype is not None else None
@property
def rank(self) -> int | None:
return None if self.shape is None else len(self.shape)
def __index__(self) -> int:
# A value only exists once the model runs (an extent along a dynamic
# axis, say); an argument the ONNX node fixes at export time cannot
# take one.
raise TypeError(
f"{self.name!r} is computed when the model runs, but this argument "
"must be a fixed integer in the exported graph"
)
def __repr__(self) -> str: # pragma: no cover - debugging aid
return f"Value({self.name!r}, shape={self.shape}, dtype={self.dtype})"
[docs]
class GraphBuilder:
"""Accumulates ONNX nodes, initializers and unique value names."""
def __init__(self, opset: int, name: str = "tensorplay_model") -> None:
self.opset = int(opset)
self.name = name
self.nodes: list[Any] = []
self.initializers: list[Any] = []
self.value_info: list[Any] = []
self._used_names: set[str] = set()
self._counters: dict[str, int] = {}
self._constant_cache: dict[tuple, str] = {}
#: Per-axis extents of a value that the exported graph fixes, ``None``
#: on the axes that vary at runtime. Unset, every sampled extent is
#: taken as fixed.
self.static_dims: Callable[[Value], Sequence[int | None] | None] | None = None
# -- naming -------------------------------------------------------------
def reserve(self, name: str) -> str:
self._used_names.add(name)
return name
def unique(self, prefix: str) -> str:
prefix = _sanitize(prefix)
index = self._counters.get(prefix, 0)
while True:
candidate = prefix if index == 0 else f"{prefix}_{index}"
index += 1
if candidate not in self._used_names:
self._counters[prefix] = index
self._used_names.add(candidate)
return candidate
# -- constants ----------------------------------------------------------
def initializer(self, array: np.ndarray, name_hint: str = "const") -> str:
name = self.unique(name_hint)
self.initializers.append(numpy_helper.from_array(array, name))
return name
[docs]
def constant(
self,
value: Any,
dtype: Any = None,
name_hint: str = "const",
) -> str:
"""Materialize a python/tensor constant as a cached initializer."""
array = _to_numpy(value)
if dtype is not None:
array = array.astype(np.dtype(dtype), copy=False)
key = (array.dtype.str, array.shape, array.tobytes())
cached = self._constant_cache.get(key)
if cached is not None:
return cached
name = self.initializer(array, name_hint)
self._constant_cache[key] = name
return name
def int64_1d(self, values: Sequence[int], name_hint: str = "axes") -> str:
return self.constant(
np.asarray(list(values), dtype=np.int64), name_hint=name_hint
)
# -- nodes --------------------------------------------------------------
[docs]
def op(
self,
op_type: str,
inputs: Sequence[str],
*,
num_outputs: int = 1,
name_hint: str | None = None,
outputs: Sequence[str] | None = None,
**attrs: Any,
) -> Any:
"""Emit one ONNX node and return its output name (or list of names)."""
hint = name_hint or op_type.lower()
if outputs is None:
outputs = [
self.unique(hint if num_outputs == 1 else f"{hint}_{index}")
for index in range(num_outputs)
]
else:
outputs = [self.reserve(name) for name in outputs]
attrs = {key: value for key, value in attrs.items() if value is not None}
self.nodes.append(
helper.make_node(op_type, list(inputs), list(outputs), **attrs)
)
return outputs[0] if len(outputs) == 1 else list(outputs)
def require_opset(self, minimum: int, feature: str) -> None:
if self.opset < minimum:
raise UnsupportedOperatorError(
f"{feature} requires ONNX opset >= {minimum}, got {self.opset}"
)
def _sanitize(name: str) -> str:
return "".join(char if char.isalnum() or char == "_" else "_" for char in str(name))
# ---------------------------------------------------------------------------
# Handler context
# ---------------------------------------------------------------------------
[docs]
class OpContext:
"""Argument access plus emission helpers handed to every handler."""
__slots__ = (
"b", "node_name", "params", "args", "kwargs", "out_shape", "out_dtype", "sample"
)
def __init__(
self,
builder: GraphBuilder,
node_name: str,
params: Sequence[str],
args: Sequence[Any],
kwargs: dict[str, Any],
out_shape: tuple | None = None,
out_dtype: np.dtype | None = None,
sample: Any = None,
) -> None:
self.b = builder
self.node_name = node_name
self.params = list(params)
self.args = list(args)
self.kwargs = dict(kwargs)
self.out_shape = out_shape
self.out_dtype = out_dtype
#: What the node produced when the graph ran on the example inputs.
self.sample = sample
# -- argument access ----------------------------------------------------
def get(self, param: str, default: Any = None) -> Any:
if param in self.kwargs:
return self.kwargs[param]
try:
index = self.params.index(param)
except ValueError as exc: # pragma: no cover - handler bug
raise KeyError(
f"{param!r} is not declared in the parameter list {self.params}"
) from exc
if index < len(self.args):
return self.args[index]
return default
def has(self, param: str) -> bool:
return self.get(param, _MISSING) is not _MISSING
@property
def x(self) -> Any:
"""First declared argument (the input tensor for nearly every op)."""
return self.get(self.params[0])
# -- value helpers ------------------------------------------------------
[docs]
def name(self, value: Any, name_hint: str = "const") -> str:
"""ONNX value name for ``value``, materializing constants on demand."""
if isinstance(value, Value):
return value.name
if isinstance(value, str):
return value
if value is None:
return ""
return self.b.constant(value, name_hint=name_hint)
[docs]
def cast_like(self, value: Any, reference: Any, name_hint: str = "const") -> str:
"""Name for ``value``, materialized with ``reference``'s dtype."""
if isinstance(value, (Value, str)):
return self.name(value)
dtype = self.dtype(reference)
return self.b.constant(value, dtype=dtype, name_hint=name_hint)
def shape(self, value: Any) -> tuple | None:
if isinstance(value, Value):
return value.shape
if value is None:
return None
try:
return tuple(_to_numpy(value).shape)
except Exception: # noqa: BLE001 - non-tensor argument
return None
def dtype(self, value: Any) -> np.dtype | None:
if isinstance(value, Value):
return value.dtype
if value is None:
return None
try:
return _to_numpy(value).dtype
except Exception: # noqa: BLE001 - non-tensor argument
return None
def rank(self, value: Any, what: str = "input") -> int:
shape = self.shape(value)
if shape is None:
raise UnsupportedOperatorError(
f"{self.node_name}: the rank of {what} is unknown; export with "
"example inputs so shapes can be propagated"
)
return len(shape)
def dim_size(self, value: Any, axis: int, what: str = "input") -> int:
shape = self.shape(value)
if shape is None:
raise UnsupportedOperatorError(
f"{self.node_name}: the shape of {what} is unknown; export with "
"example inputs so shapes can be propagated"
)
return int(shape[axis])
# -- emission -----------------------------------------------------------
def op(self, op_type: str, inputs: Sequence[Any], **kwargs: Any) -> Any:
kwargs.setdefault("name_hint", f"{self.node_name}_{op_type.lower()}")
resolved = [item if isinstance(item, str) else self.name(item) for item in inputs]
return self.b.op(op_type, resolved, **kwargs)
def unary(self, op_type: str, **attrs: Any) -> str:
return self.op(op_type, [self.x], **attrs)
def binary(self, op_type: str) -> str:
return self.op(op_type, list(_binary_operands(self, self.x, self.get(self.params[1]))))
class _Missing:
def __repr__(self) -> str: # pragma: no cover - debugging aid
return "<missing>"
_MISSING = _Missing()
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
Handler = Callable[[OpContext], Any]
_FUNCTION_HANDLERS: dict[tuple[str, str], tuple[Handler, list[str]]] = {}
_ANY_MODULE_HANDLERS: dict[str, tuple[Handler, list[str]]] = {}
_METHOD_HANDLERS: dict[str, tuple[Handler, list[str]]] = {}
_TP_MODULES = ("tensorplay.functional", "tensorplay.nn.functional", "tensorplay")
def register(
name: str,
params: str,
*,
module: str | None = None,
methods: bool = True,
) -> Callable[[Handler], Handler]:
"""Register a handler for a captured function (and same-named method)."""
param_list = params.split()
def decorate(handler: Handler) -> Handler:
entry = (handler, param_list)
if module is None:
_ANY_MODULE_HANDLERS[name] = entry
else:
_FUNCTION_HANDLERS[(module, name)] = entry
if methods:
_METHOD_HANDLERS.setdefault(name, entry)
return handler
return decorate
def register_method(name: str, params: str) -> Callable[[Handler], Handler]:
"""Register a handler used only for ``call_method`` nodes."""
param_list = params.split()
def decorate(handler: Handler) -> Handler:
_METHOD_HANDLERS[name] = (handler, param_list)
return handler
return decorate
def alias(name: str, target: str, *, params: str | None = None) -> None:
"""Register ``name`` using the handler already registered for ``target``."""
entry = _ANY_MODULE_HANDLERS.get(target) or _METHOD_HANDLERS.get(target)
if entry is None: # pragma: no cover - registration order bug
raise KeyError(f"no handler registered for {target!r}")
handler, param_list = entry
if params is not None:
param_list = params.split()
_ANY_MODULE_HANDLERS[name] = (handler, param_list)
_METHOD_HANDLERS.setdefault(name, (handler, param_list))
[docs]
def lookup_function_handler(
module: str, name: str
) -> tuple[Handler, list[str]] | None:
entry = _FUNCTION_HANDLERS.get((module, name))
if entry is not None:
return entry
return _ANY_MODULE_HANDLERS.get(name)
[docs]
def lookup_method_handler(name: str) -> tuple[Handler, list[str]] | None:
return _METHOD_HANDLERS.get(name)
# ---------------------------------------------------------------------------
# Shared lowering helpers
# ---------------------------------------------------------------------------
def _normalize_axis(axis: int, rank: int) -> int:
axis = int(axis)
return axis + rank if axis < 0 else axis
def _as_int_list(value: Any) -> list[int]:
if value is None:
return []
if isinstance(value, (list, tuple)):
return [int(item) for item in value]
return [int(value)]
def _pair_attr(value: Any, count: int) -> list[int]:
values = _as_int_list(value)
if len(values) == 1:
return values * count
if len(values) != count:
raise UnsupportedOperatorError(
f"expected {count} values, got {values}"
)
return values
def _reduce(
ctx: OpContext,
onnx_op: str,
data: Any,
dims: Any,
keepdim: bool,
*,
axes_input_since: int,
) -> str:
"""Emit a Reduce* node, honoring the opset that moved ``axes`` to an input."""
keepdims = 1 if keepdim else 0
if dims is None or (isinstance(dims, (list, tuple)) and not dims):
axes: list[int] | None = None
else:
axes = _as_int_list(dims)
if axes is None:
return ctx.op(onnx_op, [data], keepdims=keepdims)
if ctx.b.opset >= axes_input_since:
return ctx.op(
onnx_op, [data, ctx.b.int64_1d(axes, f"{ctx.node_name}_axes")], keepdims=keepdims
)
return ctx.op(onnx_op, [data], axes=axes, keepdims=keepdims)
def _reduce_sum(ctx: OpContext, data: Any, dims: Any, keepdim: bool) -> str:
return _reduce(ctx, "ReduceSum", data, dims, keepdim, axes_input_since=13)
def _squeeze(ctx: OpContext, data: Any, axes: Sequence[int]) -> str:
axes = list(axes)
if not axes:
return ctx.op("Identity", [data])
if ctx.b.opset >= 13:
return ctx.op("Squeeze", [data, ctx.b.int64_1d(axes, f"{ctx.node_name}_sq_axes")])
return ctx.op("Squeeze", [data], axes=axes)
def _unsqueeze(ctx: OpContext, data: Any, axes: Sequence[int]) -> str:
axes = list(axes)
if not axes:
return ctx.op("Identity", [data])
if ctx.b.opset >= 13:
return ctx.op("Unsqueeze", [data, ctx.b.int64_1d(axes, f"{ctx.node_name}_us_axes")])
return ctx.op("Unsqueeze", [data], axes=axes)
def _reshape(ctx: OpContext, data: Any, shape: Sequence[int]) -> str:
return ctx.op(
"Reshape",
[data, ctx.b.int64_1d(shape, f"{ctx.node_name}_shape")],
)
def _cast(ctx: OpContext, data: Any, np_dtype: Any) -> str:
return ctx.op("Cast", [data], to=int(_np_dtype_to_onnx(np_dtype)))
def _scalar(ctx: OpContext, value: Any, dtype: Any, hint: str = "scalar") -> str:
return ctx.b.constant(
np.asarray(value, dtype=np.dtype(dtype)), name_hint=f"{ctx.node_name}_{hint}"
)
def _slice(
ctx: OpContext,
data: Any,
starts: Sequence[int],
ends: Sequence[int],
axes: Sequence[int],
steps: Sequence[int] | None = None,
) -> str:
inputs = [
data,
ctx.b.int64_1d(starts, f"{ctx.node_name}_starts"),
ctx.b.int64_1d(ends, f"{ctx.node_name}_ends"),
ctx.b.int64_1d(axes, f"{ctx.node_name}_slice_axes"),
]
if steps is not None:
inputs.append(ctx.b.int64_1d(steps, f"{ctx.node_name}_steps"))
return ctx.op("Slice", inputs)
def _float_dtype(ctx: OpContext, value: Any) -> np.dtype:
dtype = ctx.dtype(value)
if dtype is None or dtype.kind != "f":
return np.dtype(np.float32)
return dtype
_KIND_ORDER = {"b": 0, "u": 1, "i": 1, "f": 2, "c": 3}
def _common_dtype(ctx: OpContext, operands: Sequence[Any]) -> np.dtype | None:
"""The dtype a binary op computes in, or ``None`` when it is not known.
ONNX requires both operands of an arithmetic node to share one type, while
the captured program mixes them freely (a float tensor times an extent,
an integer tensor divided into a float). Arithmetic computes in the dtype
it produced; a comparison computes in the widest kind among its operands,
and within that kind a tensor operand settles the width rather than a
scalar does.
"""
if ctx.out_dtype is not None and ctx.out_dtype.kind != "b":
return ctx.out_dtype
found: list[tuple[np.dtype, bool]] = []
for value in operands:
if isinstance(value, str):
return None
dtype = ctx.dtype(value)
if dtype is None:
return None
found.append((dtype, isinstance(value, Value) and bool(value.shape)))
if not found:
return None
top = max(_KIND_ORDER.get(dtype.kind, -1) for dtype, _ in found)
widest = [dtype for dtype, _ in found if _KIND_ORDER.get(dtype.kind, -1) == top]
tensors = [
dtype
for dtype, is_tensor in found
if is_tensor and _KIND_ORDER.get(dtype.kind, -1) == top
]
result = (tensors or widest)[0]
for dtype in tensors[1:]:
result = np.promote_types(result, dtype)
return np.dtype(result)
def _coerce(ctx: OpContext, value: Any, dtype: np.dtype) -> str:
"""Name for ``value`` in ``dtype``: constants are made in it, values cast."""
if isinstance(value, Value):
if value.dtype is None or value.dtype == dtype:
return value.name
return _cast(ctx, value, dtype)
if isinstance(value, str):
return value
return ctx.b.constant(value, dtype=dtype, name_hint=f"{ctx.node_name}_const")
def _binary_operands(ctx: OpContext, left: Any, right: Any) -> tuple[str, str]:
dtype = _common_dtype(ctx, (left, right))
if dtype is None:
return ctx.name(left), ctx.cast_like(right, left)
return _coerce(ctx, left, dtype), _coerce(ctx, right, dtype)
def _fixed_dims(ctx: OpContext, value: Any) -> list[int | None]:
"""Each extent of ``value`` the exported graph fixes; ``None`` where the
extent varies at runtime (it follows a dynamic axis of an input)."""
shape = ctx.shape(value)
if shape is None:
raise UnsupportedOperatorError(
f"{ctx.node_name}: the shape of the input is unknown; export with "
"example inputs so shapes can be propagated"
)
fixed: Sequence[int | None] | None = shape
if isinstance(value, Value) and ctx.b.static_dims is not None:
fixed = ctx.b.static_dims(value)
if fixed is None or len(fixed) != len(shape):
return [None] * len(shape)
return [
None if known is None else int(size) for size, known in zip(shape, fixed)
]
def _dims(ctx: OpContext, value: Any, axes: Sequence[int] | None = None) -> list[Any]:
"""The extents of ``value`` on ``axes`` (all by default), as the exported
graph knows them.
An extent the graph fixes is a plain ``int``; one that varies at runtime
is read off the tensor with ``Shape`` and arrives as a 0-d int64
:class:`Value`, so arithmetic and reshapes built on it follow the actual
input.
"""
fixed = _fixed_dims(ctx, value)
dims: list[Any] = []
shape_name = None
for axis in range(len(fixed)) if axes is None else axes:
if fixed[axis] is not None:
dims.append(fixed[axis])
continue
if shape_name is None:
shape_name = ctx.op("Shape", [value])
index = ctx.b.constant(np.asarray(axis, dtype=np.int64), name_hint="axis")
extent = ctx.op("Gather", [shape_name, index], axis=0)
dims.append(Value(extent, (), np.int64))
return dims
def _shape_tensor(ctx: OpContext, dims: Sequence[Any], hint: str = "shape") -> str:
"""A 1-D int64 shape input made from fixed ints and runtime extents."""
if not any(isinstance(item, Value) for item in dims):
return ctx.b.int64_1d([int(item) for item in dims], f"{ctx.node_name}_{hint}")
pieces: list[str] = []
run: list[int] = []
for item in dims:
if not isinstance(item, Value):
run.append(int(item))
continue
if run:
pieces.append(ctx.b.int64_1d(run, f"{ctx.node_name}_{hint}"))
run = []
extent = _coerce(ctx, item, np.dtype(np.int64))
if not item.shape:
extent = _unsqueeze(ctx, extent, [0])
pieces.append(extent)
if run:
pieces.append(ctx.b.int64_1d(run, f"{ctx.node_name}_{hint}"))
return ctx.op("Concat", pieces, axis=0)
def _is_tensor_like(value: Any) -> bool:
return hasattr(value, "shape") and hasattr(value, "numpy")
# ---------------------------------------------------------------------------
# Pointwise unary ops
# ---------------------------------------------------------------------------
_SIMPLE_UNARY = {
"abs": "Abs",
"neg": "Neg",
"exp": "Exp",
"log": "Log",
"sqrt": "Sqrt",
"ceil": "Ceil",
"floor": "Floor",
"round": "Round",
"sign": "Sign",
"sin": "Sin",
"cos": "Cos",
"tan": "Tan",
"asin": "Asin",
"acos": "Acos",
"atan": "Atan",
"sinh": "Sinh",
"cosh": "Cosh",
"asinh": "Asinh",
"acosh": "Acosh",
"atanh": "Atanh",
"erf": "Erf",
"relu": "Relu",
"sigmoid": "Sigmoid",
"tanh": "Tanh",
"reciprocal": "Reciprocal",
"logical_not": "Not",
"invert": "Not",
"isnan": "IsNaN",
"isinf": "IsInf",
"softsign": "Softsign",
"hardswish": "HardSwish",
"det": "Det",
"nonzero": "NonZero",
}
for _name, _onnx_op in _SIMPLE_UNARY.items():
register(_name, "input")(
lambda ctx, _op=_onnx_op: ctx.unary(_op)
)
_SIMPLE_BINARY = {
"mul": "Mul",
"multiply": "Mul",
"div": "Div",
"divide": "Div",
"true_divide": "Div",
"truediv": "Div",
"pow": "Pow",
"matmul": "MatMul",
"mm": "MatMul",
"bmm": "MatMul",
"maximum": "Max",
"minimum": "Min",
"eq": "Equal",
"lt": "Less",
"le": "LessOrEqual",
"gt": "Greater",
"ge": "GreaterOrEqual",
"logical_and": "And",
"logical_or": "Or",
"logical_xor": "Xor",
"and": "And",
"or": "Or",
"xor": "Xor",
"bitwise_and": "BitwiseAnd",
"bitwise_or": "BitwiseOr",
"bitwise_xor": "BitwiseXor",
}
for _name, _onnx_op in _SIMPLE_BINARY.items():
register(_name, "input other")(
lambda ctx, _op=_onnx_op: ctx.binary(_op)
)
def _scaled_operands(ctx: OpContext) -> tuple[str, str]:
"""Both operands, the second pre-multiplied by ``alpha`` when one was given."""
left, other = _binary_operands(ctx, ctx.x, ctx.get("other"))
alpha = ctx.get("alpha", 1)
if alpha is None or float(alpha) == 1.0:
return left, other
dtype = _common_dtype(ctx, (ctx.x, ctx.get("other")))
scale = (
ctx.cast_like(float(alpha), ctx.x)
if dtype is None
else _coerce(ctx, float(alpha), dtype)
)
return left, ctx.op("Mul", [other, scale])
@register("add", "input other alpha")
def _handle_add(ctx: OpContext) -> Any:
if isinstance(ctx.x, (list, tuple)):
# ``x.shape[:2] + (-1,)``: sequences of extents concatenate.
return tuple(ctx.x) + tuple(ctx.get("other"))
return ctx.op("Add", list(_scaled_operands(ctx)))
@register("sub", "input other alpha")
def _handle_sub(ctx: OpContext) -> str:
return ctx.op("Sub", list(_scaled_operands(ctx)))
alias("subtract", "sub")
@register("remainder", "input other")
def _handle_remainder(ctx: OpContext) -> str:
"""``x - floor(x / y) * y``: the result takes the divisor's sign."""
left, other = _binary_operands(ctx, ctx.x, ctx.get("other"))
dtype = _common_dtype(ctx, (ctx.x, ctx.get("other"))) or ctx.dtype(ctx.x)
if dtype is not None and dtype.kind in "iu":
# Integer Mod already rounds the quotient towards negative infinity.
return ctx.op("Mod", [left, other], fmod=0)
quotient = ctx.op("Floor", [ctx.op("Div", [left, other])])
return ctx.op("Sub", [left, ctx.op("Mul", [quotient, other])])
@register("fmod", "input other")
def _handle_fmod(ctx: OpContext) -> str:
"""``x - trunc(x / y) * y``, which is what ``Mod(fmod=1)`` computes."""
return ctx.op("Mod", list(_binary_operands(ctx, ctx.x, ctx.get("other"))), fmod=1)
@register("ne", "input other")
def _handle_ne(ctx: OpContext) -> str:
return ctx.op("Not", [ctx.binary("Equal")])
@register("floordiv", "input other")
def _handle_floordiv(ctx: OpContext) -> str:
left, other = _binary_operands(ctx, ctx.x, ctx.get("other"))
dtype = _common_dtype(ctx, (ctx.x, ctx.get("other"))) or ctx.dtype(ctx.x)
if dtype is not None and dtype.kind in "iu":
as_float = _cast(ctx, left, np.float32)
divided = ctx.op("Div", [as_float, _cast(ctx, other, np.float32)])
return _cast(ctx, ctx.op("Floor", [divided]), dtype)
return ctx.op("Floor", [ctx.op("Div", [left, other])])
alias("floor_divide", "floordiv")
@register("square", "input")
def _handle_square(ctx: OpContext) -> str:
return ctx.op("Mul", [ctx.x, ctx.x])
@register("rsqrt", "input")
def _handle_rsqrt(ctx: OpContext) -> str:
return ctx.op("Reciprocal", [ctx.op("Sqrt", [ctx.x])])
@register("log2", "input")
def _handle_log2(ctx: OpContext) -> str:
scale = _scalar(ctx, 1.0 / math.log(2.0), _float_dtype(ctx, ctx.x), "log2")
return ctx.op("Mul", [ctx.op("Log", [ctx.x]), scale])
@register("log10", "input")
def _handle_log10(ctx: OpContext) -> str:
scale = _scalar(ctx, 1.0 / math.log(10.0), _float_dtype(ctx, ctx.x), "log10")
return ctx.op("Mul", [ctx.op("Log", [ctx.x]), scale])
@register("log1p", "input")
def _handle_log1p(ctx: OpContext) -> str:
one = _scalar(ctx, 1.0, _float_dtype(ctx, ctx.x), "one")
return ctx.op("Log", [ctx.op("Add", [ctx.x, one])])
@register("expm1", "input")
def _handle_expm1(ctx: OpContext) -> str:
one = _scalar(ctx, 1.0, _float_dtype(ctx, ctx.x), "one")
return ctx.op("Sub", [ctx.op("Exp", [ctx.x]), one])
@register("clamp", "input min max")
def _handle_clamp(ctx: OpContext) -> str:
minimum = ctx.get("min")
maximum = ctx.get("max")
inputs: list[Any] = [ctx.x]
inputs.append("" if minimum is None else ctx.cast_like(minimum, ctx.x))
if maximum is not None:
inputs.append(ctx.cast_like(maximum, ctx.x))
return ctx.op("Clip", inputs)
@register("clamp_min", "input min")
def _handle_clamp_min(ctx: OpContext) -> str:
return ctx.op("Clip", [ctx.x, ctx.cast_like(ctx.get("min"), ctx.x)])
@register("clamp_max", "input max")
def _handle_clamp_max(ctx: OpContext) -> str:
return ctx.op("Clip", [ctx.x, "", ctx.cast_like(ctx.get("max"), ctx.x)])
@register("where", "condition input other")
def _handle_where(ctx: OpContext) -> str:
condition = ctx.get("condition")
body = ctx.get("input")
other = ctx.get("other")
return ctx.op(
"Where",
[condition, ctx.cast_like(body, other), ctx.cast_like(other, body)],
)
@register("masked_fill", "input mask value")
def _handle_masked_fill(ctx: OpContext) -> str:
value = ctx.get("value")
return ctx.op(
"Where", [ctx.get("mask"), ctx.cast_like(value, ctx.x), ctx.x]
)
@register("lerp", "input end weight")
def _handle_lerp(ctx: OpContext) -> str:
start, end = ctx.x, ctx.get("end")
weight = ctx.cast_like(ctx.get("weight"), start)
delta = ctx.op("Sub", [end, start])
return ctx.op("Add", [start, ctx.op("Mul", [delta, weight])])
# ---------------------------------------------------------------------------
# Activations
# ---------------------------------------------------------------------------
@register("gelu", "input approximate")
def _handle_gelu(ctx: OpContext) -> str:
approximate = ctx.get("approximate", "none") or "none"
if ctx.b.opset >= 20:
return ctx.op("Gelu", [ctx.x], approximate=str(approximate))
dtype = _float_dtype(ctx, ctx.x)
half = _scalar(ctx, 0.5, dtype, "half")
one = _scalar(ctx, 1.0, dtype, "one")
if approximate == "tanh":
alpha = _scalar(ctx, math.sqrt(2.0 / math.pi), dtype, "gelu_alpha")
beta = _scalar(ctx, 0.044715, dtype, "gelu_beta")
three = _scalar(ctx, 3.0, dtype, "three")
cubed = ctx.op("Pow", [ctx.x, three])
inner = ctx.op("Add", [ctx.x, ctx.op("Mul", [beta, cubed])])
tanh = ctx.op("Tanh", [ctx.op("Mul", [alpha, inner])])
return ctx.op(
"Mul", [ctx.op("Mul", [half, ctx.x]), ctx.op("Add", [one, tanh])]
)
inv_sqrt2 = _scalar(ctx, 1.0 / math.sqrt(2.0), dtype, "inv_sqrt2")
erf = ctx.op("Erf", [ctx.op("Mul", [ctx.x, inv_sqrt2])])
return ctx.op("Mul", [ctx.op("Mul", [half, ctx.x]), ctx.op("Add", [one, erf])])
@register("silu", "input inplace")
def _handle_silu(ctx: OpContext) -> str:
return ctx.op("Mul", [ctx.x, ctx.op("Sigmoid", [ctx.x])])
alias("swish", "silu")
@register("mish", "input")
def _handle_mish(ctx: OpContext) -> str:
if ctx.b.opset >= 18:
return ctx.op("Mish", [ctx.x])
softplus = ctx.op("Softplus", [ctx.x])
return ctx.op("Mul", [ctx.x, ctx.op("Tanh", [softplus])])
@register("hardsigmoid", "input")
def _handle_hardsigmoid(ctx: OpContext) -> str:
return ctx.op("HardSigmoid", [ctx.x], alpha=1.0 / 6.0, beta=0.5)
@register("leaky_relu", "input negative_slope inplace")
def _handle_leaky_relu(ctx: OpContext) -> str:
return ctx.op("LeakyRelu", [ctx.x], alpha=float(ctx.get("negative_slope", 0.01)))
@register("elu", "input alpha scale input_scale")
def _handle_elu(ctx: OpContext) -> str:
scale = ctx.get("scale", 1)
input_scale = ctx.get("input_scale", 1)
if float(scale) != 1.0 or float(input_scale) != 1.0:
raise UnsupportedOperatorError(
"elu with scale/input_scale != 1 has no ONNX equivalent"
)
return ctx.op("Elu", [ctx.x], alpha=float(ctx.get("alpha", 1.0)))
@register("selu", "input")
def _handle_selu(ctx: OpContext) -> str:
return ctx.op("Selu", [ctx.x])
@register("celu", "input alpha")
def _handle_celu(ctx: OpContext) -> str:
return ctx.op("Celu", [ctx.x], alpha=float(ctx.get("alpha", 1.0)))
@register("hardtanh", "input min_val max_val inplace")
def _handle_hardtanh(ctx: OpContext) -> str:
low = ctx.cast_like(float(ctx.get("min_val", -1.0)), ctx.x)
high = ctx.cast_like(float(ctx.get("max_val", 1.0)), ctx.x)
return ctx.op("Clip", [ctx.x, low, high])
@register("relu6", "input inplace")
def _handle_relu6(ctx: OpContext) -> str:
return ctx.op(
"Clip", [ctx.x, ctx.cast_like(0.0, ctx.x), ctx.cast_like(6.0, ctx.x)]
)
@register("threshold", "input threshold value inplace")
def _handle_threshold(ctx: OpContext) -> str:
limit = ctx.cast_like(float(ctx.get("threshold")), ctx.x)
value = ctx.cast_like(float(ctx.get("value")), ctx.x)
return ctx.op("Where", [ctx.op("Greater", [ctx.x, limit]), ctx.x, value])
@register("softplus", "input beta threshold")
def _handle_softplus(ctx: OpContext) -> str:
beta = float(ctx.get("beta", 1.0))
if beta == 1.0:
return ctx.op("Softplus", [ctx.x])
dtype = _float_dtype(ctx, ctx.x)
beta_const = _scalar(ctx, beta, dtype, "beta")
scaled = ctx.op("Softplus", [ctx.op("Mul", [ctx.x, beta_const])])
return ctx.op("Div", [scaled, beta_const])
@register("prelu", "input weight")
def _handle_prelu(ctx: OpContext) -> str:
weight = ctx.get("weight")
rank = ctx.rank(ctx.x)
weight_shape = ctx.shape(weight)
slope: Any = weight
if rank > 2 and weight_shape is not None and int(np.prod(weight_shape)) > 1:
slope = _reshape(
ctx, ctx.name(weight), [int(np.prod(weight_shape))] + [1] * (rank - 2)
)
return ctx.op("PRelu", [ctx.x, slope])
@register("softmax", "input dim dtype")
def _handle_softmax(ctx: OpContext) -> str:
return ctx.op("Softmax", [ctx.x], axis=int(ctx.get("dim", -1) or -1))
@register("log_softmax", "input dim dtype")
def _handle_log_softmax(ctx: OpContext) -> str:
return ctx.op("LogSoftmax", [ctx.x], axis=int(ctx.get("dim", -1) or -1))
@register("glu", "input dim")
def _handle_glu(ctx: OpContext) -> str:
dim = _normalize_axis(ctx.get("dim", -1), ctx.rank(ctx.x))
size = ctx.dim_size(ctx.x, dim)
if size % 2:
raise UnsupportedOperatorError("glu requires an even split dimension")
half = size // 2
first, second = ctx.op(
"Split",
[ctx.x, ctx.b.int64_1d([half, half], f"{ctx.node_name}_split")],
axis=dim,
num_outputs=2,
)
return ctx.op("Mul", [first, ctx.op("Sigmoid", [second])])
@register("dropout", "input p training inplace")
def _handle_dropout(ctx: OpContext) -> str:
if not bool(ctx.get("training", True)) or float(ctx.get("p", 0.5)) == 0.0:
return ctx.op("Identity", [ctx.x])
ratio = _scalar(ctx, float(ctx.get("p", 0.5)), np.float32, "ratio")
training = _scalar(ctx, True, np.bool_, "training")
outputs = ctx.op("Dropout", [ctx.x, ratio, training], num_outputs=2)
return outputs[0]
@register("identity", "input")
def _handle_identity(ctx: OpContext) -> str:
return ctx.op("Identity", [ctx.x])
for _passthrough in ("detach", "contiguous", "clone", "alias"):
register(_passthrough, "input")(lambda ctx: ctx.op("Identity", [ctx.x]))
# ---------------------------------------------------------------------------
# Reductions
# ---------------------------------------------------------------------------
def _variadic_ints(ctx: OpContext, param: str) -> list[int]:
"""Read ``x.view(2, 3)`` and ``x.view([2, 3])`` alike."""
value = ctx.get(param)
if isinstance(value, (list, tuple)):
return [int(item) for item in value]
index = ctx.params.index(param)
tail = ctx.args[index:]
if tail:
return [int(item) for item in tail]
if value is None:
return []
return [int(value)]
def _variadic_dims(ctx: OpContext, param: str) -> list[Any]:
"""Like :func:`_variadic_ints`, but runtime extents stay :class:`Value`."""
def settle(item: Any) -> Any:
return item if isinstance(item, Value) else int(item)
value = ctx.get(param)
if isinstance(value, (list, tuple)):
return [settle(item) for item in value]
index = ctx.params.index(param)
tail = ctx.args[index:]
if tail:
return [settle(item) for item in tail]
if value is None:
return []
return [settle(value)]
@register("sum", "input dim keepdim dtype")
def _handle_sum(ctx: OpContext) -> str:
return _reduce_sum(ctx, ctx.x, ctx.get("dim"), bool(ctx.get("keepdim", False)))
@register("mean", "input dim keepdim dtype")
def _handle_mean(ctx: OpContext) -> str:
return _reduce(
ctx, "ReduceMean", ctx.x, ctx.get("dim"), bool(ctx.get("keepdim", False)),
axes_input_since=18,
)
@register("prod", "input dim keepdim dtype")
def _handle_prod(ctx: OpContext) -> str:
return _reduce(
ctx, "ReduceProd", ctx.x, ctx.get("dim"), bool(ctx.get("keepdim", False)),
axes_input_since=18,
)
@register("amax", "input dim keepdim")
def _handle_amax(ctx: OpContext) -> str:
return _reduce(
ctx, "ReduceMax", ctx.x, ctx.get("dim"), bool(ctx.get("keepdim", False)),
axes_input_since=18,
)
@register("amin", "input dim keepdim")
def _handle_amin(ctx: OpContext) -> str:
return _reduce(
ctx, "ReduceMin", ctx.x, ctx.get("dim"), bool(ctx.get("keepdim", False)),
axes_input_since=18,
)
def _minmax(ctx: OpContext, reduce_op: str, arg_op: str, elementwise: str) -> Any:
"""``max``/``min``: elementwise, whole-tensor reduction, or (values, indices).
``tensorplay.functional.max(input, dim, keepdim)`` reduces; the second
argument is only a tensor for the ``Tensor.max(other)`` overload.
"""
dim = ctx.get("dim")
if isinstance(dim, Value):
return ctx.op(elementwise, [ctx.x, dim])
keepdim = bool(ctx.get("keepdim", False))
if dim is None:
return _reduce(ctx, reduce_op, ctx.x, None, keepdim, axes_input_since=18)
axis = _normalize_axis(dim, ctx.rank(ctx.x))
values = _reduce(ctx, reduce_op, ctx.x, [axis], keepdim, axes_input_since=18)
indices = ctx.op(arg_op, [ctx.x], axis=axis, keepdims=1 if keepdim else 0)
return [values, indices]
@register("max", "input dim keepdim")
def _handle_max(ctx: OpContext) -> Any:
return _minmax(ctx, "ReduceMax", "ArgMax", "Max")
@register("min", "input dim keepdim")
def _handle_min(ctx: OpContext) -> Any:
return _minmax(ctx, "ReduceMin", "ArgMin", "Min")
def _argreduce(ctx: OpContext, onnx_op: str) -> str:
dim = ctx.get("dim")
keepdim = bool(ctx.get("keepdim", False))
data: Any = ctx.x
if dim is None:
data = _reshape(ctx, ctx.x, [-1])
dim, keepdim = 0, False
return ctx.op(onnx_op, [data], axis=int(dim), keepdims=1 if keepdim else 0)
@register("argmax", "input dim keepdim")
def _handle_argmax(ctx: OpContext) -> str:
return _argreduce(ctx, "ArgMax")
@register("argmin", "input dim keepdim")
def _handle_argmin(ctx: OpContext) -> str:
return _argreduce(ctx, "ArgMin")
def _bool_reduce(ctx: OpContext, onnx_op: str) -> str:
as_int = _cast(ctx, ctx.x, np.int32)
reduced = _reduce(
ctx, onnx_op, as_int, ctx.get("dim"), bool(ctx.get("keepdim", False)),
axes_input_since=18,
)
return _cast(ctx, reduced, np.bool_)
@register("all", "input dim keepdim")
def _handle_all(ctx: OpContext) -> str:
return _bool_reduce(ctx, "ReduceMin")
@register("any", "input dim keepdim")
def _handle_any(ctx: OpContext) -> str:
return _bool_reduce(ctx, "ReduceMax")
@register("logsumexp", "input dim keepdim")
def _handle_logsumexp(ctx: OpContext) -> str:
return _reduce(
ctx, "ReduceLogSumExp", ctx.x, ctx.get("dim"),
bool(ctx.get("keepdim", False)), axes_input_since=18,
)
@register("cumsum", "input dim dtype")
def _handle_cumsum(ctx: OpContext) -> str:
axis = _scalar(ctx, int(ctx.get("dim", 0)), np.int64, "axis")
return ctx.op("CumSum", [ctx.x, axis])
def _reduced_count(ctx: OpContext, dims: Any) -> int:
shape = ctx.shape(ctx.x)
if shape is None:
raise UnsupportedOperatorError(
f"{ctx.node_name}: variance needs a known input shape"
)
if dims is None or (isinstance(dims, (list, tuple)) and not dims):
axes = range(len(shape))
else:
axes = [_normalize_axis(axis, len(shape)) for axis in _as_int_list(dims)]
count = 1
for axis in axes:
count *= int(shape[axis])
return count
def _variance(ctx: OpContext) -> tuple[str, np.dtype]:
dims = ctx.get("dim")
keepdim = bool(ctx.get("keepdim", False))
# The divisor is the element count less the correction. When none is
# given it follows ``unbiased``, which itself defaults to one.
correction = ctx.get("correction")
if correction is None:
unbiased = ctx.get("unbiased")
correction = 1 if (unbiased is None or unbiased) else 0
correction = float(correction)
dtype = _float_dtype(ctx, ctx.x)
mean = _reduce(ctx, "ReduceMean", ctx.x, dims, True, axes_input_since=18)
centered = ctx.op("Sub", [ctx.x, mean])
squares = ctx.op("Mul", [centered, centered])
total = _reduce_sum(ctx, squares, dims, keepdim)
denominator = max(_reduced_count(ctx, dims) - correction, 1.0)
return ctx.op("Div", [total, _scalar(ctx, denominator, dtype, "count")]), dtype
@register("var", "input dim unbiased keepdim correction")
def _handle_var(ctx: OpContext) -> str:
variance, _ = _variance(ctx)
return variance
@register("std", "input dim unbiased keepdim correction")
def _handle_std(ctx: OpContext) -> str:
variance, _ = _variance(ctx)
return ctx.op("Sqrt", [variance])
def _p_norm(ctx: OpContext, data: Any, p: float, dims: Any, keepdim: bool) -> str:
if p == 2.0:
return _reduce(ctx, "ReduceL2", data, dims, keepdim, axes_input_since=18)
if p == 1.0:
return _reduce(ctx, "ReduceL1", data, dims, keepdim, axes_input_since=18)
absolute = ctx.op("Abs", [data])
if math.isinf(p):
onnx_op = "ReduceMax" if p > 0 else "ReduceMin"
return _reduce(ctx, onnx_op, absolute, dims, keepdim, axes_input_since=18)
dtype = _float_dtype(ctx, data)
exponent = _scalar(ctx, p, dtype, "p")
powered = ctx.op("Pow", [absolute, exponent])
total = _reduce_sum(ctx, powered, dims, keepdim)
inverse = _scalar(ctx, 1.0 / p, dtype, "inv_p")
return ctx.op("Pow", [total, inverse])
@register("norm", "input p dim keepdim", methods=False)
def _handle_norm(ctx: OpContext) -> str:
return _p_norm(
ctx, ctx.x, float(ctx.get("p", 2.0) or 2.0), ctx.get("dim"),
bool(ctx.get("keepdim", False)),
)
@register_method("norm", "input dim p keepdim")
def _handle_norm_method(ctx: OpContext) -> str:
"""``Tensor.norm`` has a ``(p)`` and a ``(dim, p, keepdim)`` overload."""
first = ctx.get("dim")
if first is None or isinstance(first, (list, tuple)):
dims = None if first is None else _as_int_list(first)
p = float(ctx.get("p", 2.0) or 2.0)
keepdim = bool(ctx.get("keepdim", False))
else:
dims, p, keepdim = None, float(first), False
return _p_norm(ctx, ctx.x, p, dims, keepdim)
@register("topk", "input k dim largest sorted impl")
def _handle_topk(ctx: OpContext) -> list[str]:
k = ctx.b.constant(
np.asarray([int(ctx.get("k"))], dtype=np.int64), name_hint=f"{ctx.node_name}_k"
)
return ctx.op(
"TopK",
[ctx.x, k],
num_outputs=2,
axis=int(ctx.get("dim", -1)),
largest=1 if ctx.get("largest", True) else 0,
sorted=1 if ctx.get("sorted", True) else 0,
)
@register("sort", "input dim descending")
def _handle_sort(ctx: OpContext) -> list[str]:
rank = ctx.rank(ctx.x)
axis = _normalize_axis(ctx.get("dim", -1), rank)
size = ctx.dim_size(ctx.x, axis)
k = ctx.b.constant(
np.asarray([size], dtype=np.int64), name_hint=f"{ctx.node_name}_k"
)
return ctx.op(
"TopK",
[ctx.x, k],
num_outputs=2,
axis=axis,
largest=1 if ctx.get("descending", False) else 0,
sorted=1,
)
# ---------------------------------------------------------------------------
# Shape / layout
# ---------------------------------------------------------------------------
@register("reshape", "input shape")
def _handle_reshape(ctx: OpContext) -> str:
return ctx.op("Reshape", [ctx.x, _shape_tensor(ctx, _variadic_dims(ctx, "shape"))])
alias("view", "reshape", params="input shape")
def _leading(fixed: Sequence[int | None]) -> list[int]:
"""Extents a ``Reshape`` keeps in place: one that varies is copied (0)."""
return [0 if size is None else size for size in fixed]
@register("flatten", "input start_dim end_dim")
def _handle_flatten(ctx: OpContext) -> str:
"""ONNX ``Flatten`` always yields a 2-D tensor, so only the ``start_dim=1``
case maps onto it directly; every other range becomes a ``Reshape``.
"""
rank = ctx.rank(ctx.x)
start = _normalize_axis(ctx.get("start_dim", 0) or 0, rank)
end_dim = ctx.get("end_dim", -1)
end = _normalize_axis(-1 if end_dim is None else end_dim, rank)
if end == rank - 1:
if start == 0:
return _reshape(ctx, ctx.x, [-1])
if start == 1:
return ctx.op("Flatten", [ctx.x], axis=1)
fixed = _fixed_dims(ctx, ctx.x)
span = fixed[start : end + 1]
if None in span:
merged = -1
else:
merged = int(np.prod(span)) if end >= start else 1
new_shape = (
_leading(fixed[:start])
+ [merged]
+ _dims(ctx, ctx.x, range(end + 1, rank))
)
return ctx.op("Reshape", [ctx.x, _shape_tensor(ctx, new_shape)])
@register("unflatten", "input dim sizes")
def _handle_unflatten(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x)
dim = _normalize_axis(ctx.get("dim"), rank)
fixed = _fixed_dims(ctx, ctx.x)
sizes = _as_int_list(ctx.get("sizes"))
if -1 in sizes and fixed[dim] is not None:
known = int(np.prod([size for size in sizes if size != -1])) or 1
sizes = [fixed[dim] // known if size == -1 else size for size in sizes]
new_shape = (
_leading(fixed[:dim]) + sizes + _dims(ctx, ctx.x, range(dim + 1, rank))
)
return ctx.op("Reshape", [ctx.x, _shape_tensor(ctx, new_shape)])
@register("transpose", "input dim0 dim1")
def _handle_transpose(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x)
dim0 = _normalize_axis(ctx.get("dim0", 0), rank)
dim1 = _normalize_axis(ctx.get("dim1", 1), rank)
perm = list(range(rank))
perm[dim0], perm[dim1] = perm[dim1], perm[dim0]
return ctx.op("Transpose", [ctx.x], perm=perm)
alias("swapaxes", "transpose")
alias("swapdims", "transpose")
@register("t", "input")
def _handle_t(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x)
if rank < 2:
return ctx.op("Identity", [ctx.x])
return ctx.op("Transpose", [ctx.x], perm=[1, 0])
@register("permute", "input dims")
def _handle_permute(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x)
perm = [_normalize_axis(axis, rank) for axis in _variadic_ints(ctx, "dims")]
return ctx.op("Transpose", [ctx.x], perm=perm)
@register("squeeze", "input dim")
def _handle_squeeze(ctx: OpContext) -> str:
dim = ctx.get("dim")
if dim is None:
shape = ctx.shape(ctx.x)
if shape is None:
return ctx.op("Squeeze", [ctx.x])
axes = [index for index, size in enumerate(shape) if size == 1]
return _squeeze(ctx, ctx.x, axes)
rank = ctx.rank(ctx.x)
return _squeeze(ctx, ctx.x, [_normalize_axis(axis, rank) for axis in _as_int_list(dim)])
@register("unsqueeze", "input dim")
def _handle_unsqueeze(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x) + 1
axes = [_normalize_axis(axis, rank) for axis in _as_int_list(ctx.get("dim"))]
return _unsqueeze(ctx, ctx.x, axes)
@register("expand", "input size implicit")
def _handle_expand(ctx: OpContext) -> str:
# ``Expand`` broadcasts both ways, so an extent of 1 keeps the input's own
# extent on that axis: ``-1`` becomes 1 and holds for any input size.
sizes = [
1 if not isinstance(size, Value) and size == -1 else size
for size in _variadic_dims(ctx, "size")
]
return ctx.op("Expand", [ctx.x, _shape_tensor(ctx, sizes)])
alias("broadcast_to", "expand", params="input size")
@register("expand_as", "input other")
def _handle_expand_as(ctx: OpContext) -> str:
other = ctx.get("other")
return ctx.op("Expand", [ctx.x, ctx.op("Shape", [other])])
@register("repeat", "input repeats")
def _handle_repeat(ctx: OpContext) -> str:
repeats = _variadic_ints(ctx, "repeats")
rank = ctx.rank(ctx.x)
data: Any = ctx.x
if len(repeats) > rank:
data = _unsqueeze(ctx, ctx.x, list(range(len(repeats) - rank)))
return ctx.op(
"Tile", [data, ctx.b.int64_1d(repeats, f"{ctx.node_name}_repeats")]
)
alias("tile", "repeat", params="input dims")
@register("cat", "tensors dim")
def _handle_cat(ctx: OpContext) -> str:
tensors = ctx.get("tensors")
if not isinstance(tensors, (list, tuple)):
tensors = [tensors]
return ctx.op("Concat", list(tensors), axis=int(ctx.get("dim", 0) or 0))
alias("concat", "cat")
alias("concatenate", "cat")
@register("stack", "tensors dim")
def _handle_stack(ctx: OpContext) -> str:
tensors = ctx.get("tensors")
if not isinstance(tensors, (list, tuple)):
tensors = [tensors]
axis = int(ctx.get("dim", 0) or 0)
rank = ctx.rank(tensors[0]) + 1
axis = _normalize_axis(axis, rank)
expanded = [_unsqueeze(ctx, tensor, [axis]) for tensor in tensors]
return ctx.op("Concat", expanded, axis=axis)
def _emit_split(ctx: OpContext, data: Any, sizes: Sequence[int], axis: int) -> list[str]:
return ctx.op(
"Split",
[data, ctx.b.int64_1d(sizes, f"{ctx.node_name}_split")],
axis=axis,
num_outputs=len(sizes),
)
@register("split", "input split_size dim")
def _handle_split(ctx: OpContext) -> Any:
rank = ctx.rank(ctx.x)
axis = _normalize_axis(ctx.get("dim", 0) or 0, rank)
total = ctx.dim_size(ctx.x, axis)
split_size = ctx.get("split_size")
if isinstance(split_size, (list, tuple)):
sizes = [int(size) for size in split_size]
else:
step = int(split_size)
sizes = [step] * (total // step)
if total % step:
sizes.append(total % step)
if len(sizes) == 1:
# Still a one-element sequence, so downstream indexing keeps working.
return [ctx.op("Identity", [ctx.x])]
return _emit_split(ctx, ctx.x, sizes, axis)
@register("chunk", "input chunks dim")
def _handle_chunk(ctx: OpContext) -> Any:
rank = ctx.rank(ctx.x)
axis = _normalize_axis(ctx.get("dim", 0) or 0, rank)
total = ctx.dim_size(ctx.x, axis)
chunks = int(ctx.get("chunks"))
step = -(-total // chunks)
sizes = []
remaining = total
while remaining > 0:
sizes.append(min(step, remaining))
remaining -= sizes[-1]
if len(sizes) == 1:
return [ctx.op("Identity", [ctx.x])]
return _emit_split(ctx, ctx.x, sizes, axis)
@register("narrow", "input dim start length")
def _handle_narrow(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x)
axis = _normalize_axis(ctx.get("dim"), rank)
start = int(ctx.get("start"))
length = int(ctx.get("length"))
return _slice(ctx, ctx.x, [start], [start + length], [axis])
@register("flip", "input dims")
def _handle_flip(ctx: OpContext) -> str:
rank = ctx.rank(ctx.x)
axes = [_normalize_axis(axis, rank) for axis in _variadic_ints(ctx, "dims")]
return _slice(
ctx, ctx.x, [-1] * len(axes), [_INT64_MIN] * len(axes), axes, [-1] * len(axes)
)
@register("gather", "input dim index")
def _handle_gather(ctx: OpContext) -> str:
index = _cast(ctx, ctx.get("index"), np.int64)
return ctx.op("GatherElements", [ctx.x, index], axis=int(ctx.get("dim")))
@register("index_select", "input dim index")
def _handle_index_select(ctx: OpContext) -> str:
index = _cast(ctx, ctx.get("index"), np.int64)
return ctx.op("Gather", [ctx.x, index], axis=int(ctx.get("dim")))
@register("take", "input index")
def _handle_take(ctx: OpContext) -> str:
flat = _reshape(ctx, ctx.x, [-1])
index = _cast(ctx, ctx.get("index"), np.int64)
return ctx.op("Gather", [flat, index], axis=0)
@register("tril", "input diagonal")
def _handle_tril(ctx: OpContext) -> str:
k = _scalar(ctx, int(ctx.get("diagonal", 0) or 0), np.int64, "k")
return ctx.op("Trilu", [ctx.x, k], upper=0)
@register("triu", "input diagonal")
def _handle_triu(ctx: OpContext) -> str:
k = _scalar(ctx, int(ctx.get("diagonal", 0) or 0), np.int64, "k")
return ctx.op("Trilu", [ctx.x, k], upper=1)
@register("one_hot", "input num_classes")
def _handle_one_hot(ctx: OpContext) -> str:
num_classes = int(ctx.get("num_classes", -1))
if num_classes < 0:
raise UnsupportedOperatorError(
"one_hot needs an explicit num_classes to export (num_classes=-1 "
"depends on the runtime values of the input)"
)
depth = _scalar(ctx, num_classes, np.int64, "depth")
values = ctx.b.constant(
np.asarray([0, 1], dtype=np.int64), name_hint=f"{ctx.node_name}_values"
)
return ctx.op("OneHot", [_cast(ctx, ctx.x, np.int64), depth, values], axis=-1)
@register("pixel_shuffle", "input upscale_factor")
def _handle_pixel_shuffle(ctx: OpContext) -> str:
return ctx.op(
"DepthToSpace", [ctx.x], blocksize=int(ctx.get("upscale_factor")), mode="CRD"
)
@register("pixel_unshuffle", "input downscale_factor")
def _handle_pixel_unshuffle(ctx: OpContext) -> str:
# ONNX SpaceToDepth interleaves the channel axis in the opposite order
# (DCR), so the required layout is spelled out explicitly.
factor = int(ctx.get("downscale_factor"))
batch, channels, height, width = _fixed_dims(ctx, ctx.x)
(batch,) = _leading([batch])
if None in (channels, height, width):
raise UnsupportedOperatorError(
f"{ctx.node_name}: pixel_unshuffle exports only with fixed channel "
"and spatial extents"
)
split = _reshape(
ctx,
ctx.x,
[batch, channels, height // factor, factor, width // factor, factor],
)
ordered = ctx.op("Transpose", [split], perm=[0, 1, 3, 5, 2, 4])
return _reshape(
ctx,
ordered,
[batch, channels * factor * factor, height // factor, width // factor],
)
_CAST_METHODS = {
"float": np.float32,
"double": np.float64,
"half": np.float16,
"long": np.int64,
"int": np.int32,
"short": np.int16,
"char": np.int8,
"byte": np.uint8,
"bool": np.bool_,
}
for _method, _np_type in _CAST_METHODS.items():
register_method(_method, "input")(
lambda ctx, _type=_np_type: _cast(ctx, ctx.x, _type)
)
@register("to", "input dtype")
def _handle_to(ctx: OpContext) -> str:
dtype = ctx.get("dtype")
if dtype is None or isinstance(dtype, str):
# ``.to(device)`` / ``.to("cpu")`` is a no-op for the exported graph.
return ctx.op("Identity", [ctx.x])
try:
onnx_type = _dtype_to_onnx(dtype)
except TypeError:
return ctx.op("Identity", [ctx.x])
return ctx.op("Cast", [ctx.x], to=int(onnx_type))
alias("type", "to")
@register("type_as", "input other")
def _handle_type_as(ctx: OpContext) -> str:
dtype = ctx.dtype(ctx.get("other"))
if dtype is None:
return ctx.op("Identity", [ctx.x])
return _cast(ctx, ctx.x, dtype)
# ---------------------------------------------------------------------------
# Linear algebra
# ---------------------------------------------------------------------------
@register("linear", "input weight bias")
def _handle_linear(ctx: OpContext) -> str:
weight, bias = ctx.get("weight"), ctx.get("bias")
shape = ctx.shape(ctx.x)
if shape is not None and len(shape) == 2:
inputs = [ctx.x, weight] + ([bias] if bias is not None else [])
return ctx.op("Gemm", inputs, transB=1)
transposed = ctx.op("Transpose", [weight], perm=[1, 0])
product = ctx.op("MatMul", [ctx.x, transposed])
if bias is None:
return product
return ctx.op("Add", [product, bias])
@register("addmm", "input mat1 mat2 beta alpha")
def _handle_addmm(ctx: OpContext) -> str:
return ctx.op(
"Gemm",
[ctx.get("mat1"), ctx.get("mat2"), ctx.x],
alpha=float(ctx.get("alpha", 1)),
beta=float(ctx.get("beta", 1)),
)
@register("baddbmm", "input batch1 batch2 beta alpha")
def _handle_baddbmm(ctx: OpContext) -> str:
product = ctx.op("MatMul", [ctx.get("batch1"), ctx.get("batch2")])
alpha = float(ctx.get("alpha", 1))
beta = float(ctx.get("beta", 1))
dtype = _float_dtype(ctx, ctx.x)
if alpha != 1.0:
product = ctx.op("Mul", [product, _scalar(ctx, alpha, dtype, "alpha")])
base: Any = ctx.x
if beta != 1.0:
base = ctx.op("Mul", [ctx.x, _scalar(ctx, beta, dtype, "beta")])
return ctx.op("Add", [base, product])
# ---------------------------------------------------------------------------
# Convolution and pooling
# ---------------------------------------------------------------------------
def _spatial_rank(ctx: OpContext, weight: Any) -> int:
shape = ctx.shape(weight)
if shape is None:
raise UnsupportedOperatorError(
f"{ctx.node_name}: convolution needs a known weight shape"
)
return len(shape) - 2
def _conv_attrs(ctx: OpContext, spatial: int) -> dict[str, Any]:
padding = ctx.get("padding", 0)
attrs: dict[str, Any] = {
"strides": _pair_attr(ctx.get("stride", 1), spatial),
"dilations": _pair_attr(ctx.get("dilation", 1), spatial),
"group": int(ctx.get("groups", 1)),
}
if isinstance(padding, str):
attrs["auto_pad"] = "SAME_UPPER" if padding == "same" else "VALID"
else:
pads = _pair_attr(padding, spatial)
attrs["pads"] = pads + pads
return attrs
@register("conv1d", "input weight bias stride padding dilation groups")
@register("conv2d", "input weight bias stride padding dilation groups")
@register("conv3d", "input weight bias stride padding dilation groups")
def _handle_conv(ctx: OpContext) -> str:
weight = ctx.get("weight")
bias = ctx.get("bias")
spatial = _spatial_rank(ctx, weight)
inputs = [ctx.x, weight] + ([bias] if bias is not None else [])
attrs = _conv_attrs(ctx, spatial)
attrs["kernel_shape"] = list(ctx.shape(weight)[2:])
return ctx.op("Conv", inputs, **attrs)
@register("conv_transpose1d", "input weight bias stride padding output_padding groups dilation")
@register("conv_transpose2d", "input weight bias stride padding output_padding groups dilation")
@register("conv_transpose3d", "input weight bias stride padding output_padding groups dilation")
def _handle_conv_transpose(ctx: OpContext) -> str:
weight = ctx.get("weight")
bias = ctx.get("bias")
spatial = _spatial_rank(ctx, weight)
inputs = [ctx.x, weight] + ([bias] if bias is not None else [])
attrs = _conv_attrs(ctx, spatial)
attrs["kernel_shape"] = list(ctx.shape(weight)[2:])
output_padding = _pair_attr(ctx.get("output_padding", 0), spatial)
if any(output_padding):
attrs["output_padding"] = output_padding
return ctx.op("ConvTranspose", inputs, **attrs)
def _pool_attrs(ctx: OpContext, spatial: int) -> dict[str, Any]:
pads = _pair_attr(ctx.get("padding", 0), spatial)
return {
"kernel_shape": _pair_attr(ctx.get("kernel_size"), spatial),
"strides": _pair_attr(
ctx.get("stride") if ctx.get("stride") is not None else ctx.get("kernel_size"),
spatial,
),
"pads": pads + pads,
"ceil_mode": 1 if ctx.get("ceil_mode", False) else 0,
}
@register("max_pool1d", "input kernel_size stride padding dilation ceil_mode return_indices")
@register("max_pool2d", "input kernel_size stride padding dilation ceil_mode return_indices")
@register("max_pool3d", "input kernel_size stride padding dilation ceil_mode return_indices")
def _handle_max_pool(ctx: OpContext) -> Any:
spatial = ctx.rank(ctx.x) - 2
attrs = _pool_attrs(ctx, spatial)
attrs["dilations"] = _pair_attr(ctx.get("dilation", 1), spatial)
if not ctx.get("return_indices", False):
return ctx.op("MaxPool", [ctx.x], **attrs)
values, indices = ctx.op("MaxPool", [ctx.x], num_outputs=2, **attrs)
# ONNX numbers the indices flat over (C, *spatial); this op numbers them
# within each (batch, channel) plane, so the plane origin is subtracted.
flat_attrs = {
"kernel_shape": [1] * spatial,
"strides": [1] * spatial,
"pads": [0] * (2 * spatial),
}
_, plane_indices = ctx.op("MaxPool", [ctx.x], num_outputs=2, **flat_attrs)
axes = list(range(2, 2 + spatial))
origin = _slice(ctx, plane_indices, [0] * spatial, [1] * spatial, axes)
return [values, ctx.op("Sub", [indices, origin])]
@register("avg_pool1d", "input kernel_size stride padding ceil_mode count_include_pad divisor_override")
@register("avg_pool2d", "input kernel_size stride padding ceil_mode count_include_pad divisor_override")
@register("avg_pool3d", "input kernel_size stride padding ceil_mode count_include_pad divisor_override")
def _handle_avg_pool(ctx: OpContext) -> str:
if ctx.get("divisor_override") is not None:
raise UnsupportedOperatorError(
"avg_pool with divisor_override has no ONNX equivalent"
)
spatial = ctx.rank(ctx.x) - 2
attrs = _pool_attrs(ctx, spatial)
attrs["count_include_pad"] = 1 if ctx.get("count_include_pad", True) else 0
return ctx.op("AveragePool", [ctx.x], **attrs)
def _adaptive_pool(ctx: OpContext, global_op: str, pool_op: str) -> str:
spatial = ctx.rank(ctx.x) - 2
output_size = _pair_attr(ctx.get("output_size"), spatial)
if all(size == 1 for size in output_size):
return ctx.op(global_op, [ctx.x])
shape = ctx.shape(ctx.x)
spatial_shape = [int(size) for size in shape[2:]]
if any(
size % out for size, out in zip(spatial_shape, output_size)
): # pragma: no cover - only evenly tiled windows are representable
raise UnsupportedOperatorError(
"adaptive pooling only exports when the input divides the output size "
f"evenly (input {spatial_shape}, output {output_size})"
)
strides = [size // out for size, out in zip(spatial_shape, output_size)]
return ctx.op(
pool_op,
[ctx.x],
kernel_shape=strides,
strides=strides,
pads=[0] * (2 * spatial),
)
@register("adaptive_avg_pool1d", "input output_size")
@register("adaptive_avg_pool2d", "input output_size")
@register("adaptive_avg_pool3d", "input output_size")
def _handle_adaptive_avg_pool(ctx: OpContext) -> str:
return _adaptive_pool(ctx, "GlobalAveragePool", "AveragePool")
@register("adaptive_max_pool1d", "input output_size")
@register("adaptive_max_pool2d", "input output_size")
@register("adaptive_max_pool3d", "input output_size")
def _handle_adaptive_max_pool(ctx: OpContext) -> str:
return _adaptive_pool(ctx, "GlobalMaxPool", "MaxPool")
# ---------------------------------------------------------------------------
# Normalization
# ---------------------------------------------------------------------------
def _affine_or_default(ctx: OpContext, value: Any, size: int, fill: float, hint: str) -> Any:
if value is not None:
return value
dtype = _float_dtype(ctx, ctx.x)
return ctx.b.constant(
np.full((size,), fill, dtype=dtype), name_hint=f"{ctx.node_name}_{hint}"
)
@register("batch_norm", "input running_mean running_var weight bias training momentum eps")
def _handle_batch_norm(ctx: OpContext) -> str:
channels = ctx.dim_size(ctx.x, 1)
scale = _affine_or_default(ctx, ctx.get("weight"), channels, 1.0, "scale")
offset = _affine_or_default(ctx, ctx.get("bias"), channels, 0.0, "offset")
mean = _affine_or_default(ctx, ctx.get("running_mean"), channels, 0.0, "mean")
variance = _affine_or_default(ctx, ctx.get("running_var"), channels, 1.0, "var")
epsilon = float(ctx.get("eps", 1e-5))
if ctx.get("training", False):
ctx.b.require_opset(14, "batch_norm in training mode")
outputs = ctx.op(
"BatchNormalization",
[ctx.x, scale, offset, mean, variance],
num_outputs=3,
epsilon=epsilon,
momentum=1.0 - float(ctx.get("momentum", 0.1)),
training_mode=1,
)
return outputs[0]
return ctx.op(
"BatchNormalization",
[ctx.x, scale, offset, mean, variance],
epsilon=epsilon,
momentum=1.0 - float(ctx.get("momentum", 0.1)),
)
@register("layer_norm", "input normalized_shape weight bias eps")
def _handle_layer_norm(ctx: OpContext) -> str:
normalized_shape = _as_int_list(ctx.get("normalized_shape"))
axis = -len(normalized_shape)
size = int(np.prod(normalized_shape))
epsilon = float(ctx.get("eps", 1e-5))
weight = ctx.get("weight")
bias = ctx.get("bias")
if ctx.b.opset >= 17:
scale = weight if weight is not None else ctx.b.constant(
np.ones(normalized_shape, dtype=_float_dtype(ctx, ctx.x)),
name_hint=f"{ctx.node_name}_scale",
)
inputs = [ctx.x, scale] + ([bias] if bias is not None else [])
outputs = ctx.op(
"LayerNormalization", inputs, num_outputs=1, axis=axis, epsilon=epsilon
)
return outputs
dtype = _float_dtype(ctx, ctx.x)
axes = list(range(axis, 0))
mean = _reduce(ctx, "ReduceMean", ctx.x, axes, True, axes_input_since=18)
centered = ctx.op("Sub", [ctx.x, mean])
variance = _reduce(
ctx, "ReduceMean", ctx.op("Mul", [centered, centered]), axes, True,
axes_input_since=18,
)
denominator = ctx.op(
"Sqrt", [ctx.op("Add", [variance, _scalar(ctx, epsilon, dtype, "eps")])]
)
result: Any = ctx.op("Div", [centered, denominator])
if weight is not None:
result = ctx.op("Mul", [result, weight])
if bias is not None:
result = ctx.op("Add", [result, bias])
return result
@register("group_norm", "input num_groups weight bias eps")
def _handle_group_norm(ctx: OpContext) -> str:
groups = int(ctx.get("num_groups"))
rank = ctx.rank(ctx.x)
channels = ctx.dim_size(ctx.x, 1)
dtype = _float_dtype(ctx, ctx.x)
epsilon = float(ctx.get("eps", 1e-5))
original_shape = ctx.op("Shape", [ctx.x])
grouped = _reshape(ctx, ctx.x, [0, groups, -1])
mean = _reduce(ctx, "ReduceMean", grouped, [2], True, axes_input_since=18)
centered = ctx.op("Sub", [grouped, mean])
variance = _reduce(
ctx, "ReduceMean", ctx.op("Mul", [centered, centered]), [2], True,
axes_input_since=18,
)
denominator = ctx.op(
"Sqrt", [ctx.op("Add", [variance, _scalar(ctx, epsilon, dtype, "eps")])]
)
normalized = ctx.op("Div", [centered, denominator])
result: Any = ctx.op("Reshape", [normalized, original_shape])
affine_shape = [channels] + [1] * (rank - 2)
weight, bias = ctx.get("weight"), ctx.get("bias")
if weight is not None:
result = ctx.op("Mul", [result, _reshape(ctx, weight, affine_shape)])
if bias is not None:
result = ctx.op("Add", [result, _reshape(ctx, bias, affine_shape)])
return result
@register(
"instance_norm",
"input running_mean running_var weight bias use_input_stats momentum eps",
)
def _handle_instance_norm(ctx: OpContext) -> str:
if not ctx.get("use_input_stats", True):
raise UnsupportedOperatorError(
"instance_norm with running statistics has no ONNX equivalent"
)
channels = ctx.dim_size(ctx.x, 1)
scale = _affine_or_default(ctx, ctx.get("weight"), channels, 1.0, "scale")
offset = _affine_or_default(ctx, ctx.get("bias"), channels, 0.0, "offset")
return ctx.op(
"InstanceNormalization",
[ctx.x, scale, offset],
epsilon=float(ctx.get("eps", 1e-5)),
)
@register("local_response_norm", "input size alpha beta k")
def _handle_local_response_norm(ctx: OpContext) -> str:
size = int(ctx.get("size"))
if size % 2 == 0:
raise UnsupportedOperatorError(
f"ONNX LRN requires an odd window size, got {size}"
)
# ONNX LRN already divides the window sum by ``size``, so alpha passes
# through unscaled.
return ctx.op(
"LRN",
[ctx.x],
size=size,
alpha=float(ctx.get("alpha", 1e-4)),
beta=float(ctx.get("beta", 0.75)),
bias=float(ctx.get("k", 1.0)),
)
@register("embedding", "input weight padding_idx max_norm norm_type scale_grad_by_freq sparse")
def _handle_embedding(ctx: OpContext) -> str:
if ctx.get("max_norm") is not None:
raise UnsupportedOperatorError("embedding with max_norm has no ONNX equivalent")
indices = _cast(ctx, ctx.x, np.int64)
return ctx.op("Gather", [ctx.get("weight"), indices], axis=0)
@register("normalize", "input p dim eps")
def _handle_normalize(ctx: OpContext) -> str:
dim = int(ctx.get("dim", 1))
p = float(ctx.get("p", 2.0))
eps = float(ctx.get("eps", 1e-12))
denominator = _p_norm(ctx, ctx.x, p, [dim], True)
clipped = ctx.op(
"Clip", [denominator, _scalar(ctx, eps, _float_dtype(ctx, ctx.x), "eps")]
)
return ctx.op("Div", [ctx.x, clipped])
# ---------------------------------------------------------------------------
# Resizing and padding
# ---------------------------------------------------------------------------
_RESIZE_MODES = {
"nearest": "nearest",
"nearest-exact": "nearest",
"linear": "linear",
"bilinear": "linear",
"trilinear": "linear",
"bicubic": "cubic",
"area": "linear",
}
@register(
"interpolate",
"input size scale_factor mode align_corners recompute_scale_factor antialias",
)
def _handle_interpolate(ctx: OpContext) -> str:
mode = str(ctx.get("mode", "nearest") or "nearest")
if mode not in _RESIZE_MODES:
raise UnsupportedOperatorError(f"interpolate mode {mode!r} is not supported")
if mode == "area":
raise UnsupportedOperatorError("interpolate(mode='area') has no ONNX equivalent")
align_corners = bool(ctx.get("align_corners") or False)
rank = ctx.rank(ctx.x)
spatial = rank - 2
if align_corners:
coordinate_mode = "align_corners"
elif mode.startswith("nearest"):
coordinate_mode = "asymmetric"
else:
coordinate_mode = "tensorplay_half_pixel"
attrs: dict[str, Any] = {
"mode": _RESIZE_MODES[mode],
"coordinate_transformation_mode": coordinate_mode,
}
if _RESIZE_MODES[mode] == "nearest":
attrs["nearest_mode"] = "floor"
size = ctx.get("size")
scale_factor = ctx.get("scale_factor")
empty = ctx.b.constant(np.zeros((0,), dtype=np.float32), name_hint="resize_roi")
if size is not None:
sizes = _pair_attr(size, spatial)
sizes_name = _shape_tensor(ctx, _dims(ctx, ctx.x)[:2] + sizes, "sizes")
empty_scales = ctx.b.constant(
np.zeros((0,), dtype=np.float32), name_hint="resize_scales"
)
return ctx.op("Resize", [ctx.x, empty, empty_scales, sizes_name], **attrs)
if scale_factor is None:
raise UnsupportedOperatorError("interpolate needs size or scale_factor")
raw = list(scale_factor) if isinstance(scale_factor, (list, tuple)) else [scale_factor]
if len(raw) == 1:
raw = raw * spatial
if len(raw) != spatial:
raise UnsupportedOperatorError(
f"interpolate expects {spatial} scale factors, got {raw}"
)
factors = [float(item) for item in raw]
scales = ctx.b.constant(
np.asarray([1.0, 1.0] + factors, dtype=np.float32),
name_hint=f"{ctx.node_name}_scales",
)
return ctx.op("Resize", [ctx.x, empty, scales], **attrs)
_PAD_MODES = {
"constant": "constant",
"reflect": "reflect",
"replicate": "edge",
"circular": "wrap",
}
@register("pad", "input pad mode value")
def _handle_pad(ctx: OpContext) -> str:
mode = str(ctx.get("mode", "constant") or "constant")
if mode not in _PAD_MODES:
raise UnsupportedOperatorError(f"pad mode {mode!r} is not supported")
if mode == "circular":
ctx.b.require_opset(19, "pad(mode='circular')")
amounts = _as_int_list(ctx.get("pad"))
if len(amounts) % 2:
raise UnsupportedOperatorError("pad expects pairs of (begin, end) values")
rank = ctx.rank(ctx.x)
begins = [0] * rank
ends = [0] * rank
# The padding list starts at the LAST dimension, in (begin, end) pairs.
for index in range(len(amounts) // 2):
axis = rank - 1 - index
begins[axis] = amounts[2 * index]
ends[axis] = amounts[2 * index + 1]
pads = ctx.b.int64_1d(begins + ends, f"{ctx.node_name}_pads")
inputs: list[Any] = [ctx.x, pads]
value = ctx.get("value", 0)
if mode == "constant" and value is not None:
inputs.append(ctx.cast_like(float(value), ctx.x))
return ctx.op("Pad", inputs, mode=_PAD_MODES[mode])
# ---------------------------------------------------------------------------
# Losses
# ---------------------------------------------------------------------------
def _apply_reduction(ctx: OpContext, value: str, reduction: str) -> str:
if reduction == "none":
return value
if reduction == "sum":
return _reduce_sum(ctx, value, None, False)
return _reduce(ctx, "ReduceMean", value, None, False, axes_input_since=18)
@register("mse_loss", "input target reduction")
def _handle_mse_loss(ctx: OpContext) -> str:
difference = ctx.op("Sub", [ctx.x, ctx.get("target")])
squares = ctx.op("Mul", [difference, difference])
return _apply_reduction(ctx, squares, str(ctx.get("reduction", "mean") or "mean"))
@register("l1_loss", "input target size_average reduce reduction weight")
def _handle_l1_loss(ctx: OpContext) -> str:
difference = ctx.op("Abs", [ctx.op("Sub", [ctx.x, ctx.get("target")])])
return _apply_reduction(ctx, difference, str(ctx.get("reduction", "mean") or "mean"))
@register(
"cross_entropy",
"input target weight size_average ignore_index reduce reduction label_smoothing",
)
def _handle_cross_entropy(ctx: OpContext) -> str:
if float(ctx.get("label_smoothing", 0.0) or 0.0) != 0.0:
raise UnsupportedOperatorError(
"cross_entropy with label_smoothing has no ONNX equivalent"
)
target = ctx.get("target")
dtype = ctx.dtype(target)
if dtype is not None and dtype.kind == "f":
raise UnsupportedOperatorError(
"cross_entropy with probability targets has no ONNX equivalent"
)
inputs: list[Any] = [ctx.x, _cast(ctx, target, np.int64)]
weight = ctx.get("weight")
if weight is not None:
inputs.append(weight)
return ctx.op(
"SoftmaxCrossEntropyLoss",
inputs,
reduction=str(ctx.get("reduction", "mean") or "mean"),
ignore_index=int(ctx.get("ignore_index", -100)),
)
@register("nll_loss", "input target weight size_average ignore_index reduce reduction")
def _handle_nll_loss(ctx: OpContext) -> str:
inputs: list[Any] = [ctx.x, _cast(ctx, ctx.get("target"), np.int64)]
weight = ctx.get("weight")
if weight is not None:
inputs.append(weight)
return ctx.op(
"NegativeLogLikelihoodLoss",
inputs,
reduction=str(ctx.get("reduction", "mean") or "mean"),
ignore_index=int(ctx.get("ignore_index", -100)),
)
# ---------------------------------------------------------------------------
# Random factories
# ---------------------------------------------------------------------------
#
# A factory that draws its values is a call in the model, as it is in the
# program: every run draws again. One whose extent is written out takes it as
# an attribute; one shaped after another value follows that value at run time.
def _drawn(ctx: OpContext, op_type: str, **attrs: Any) -> str:
if ctx.out_shape is None or ctx.out_dtype is None:
raise UnsupportedOperatorError(
f"{op_type} needs the extent and element type of what it draws, "
"and this call was captured without them"
)
return ctx.op(
op_type,
[],
dtype=int(_np_dtype_to_onnx(ctx.out_dtype)),
shape=[int(extent) for extent in ctx.out_shape],
**attrs,
)
def _drawn_like(ctx: OpContext, op_type: str, **attrs: Any) -> str:
dtype = None
if ctx.out_dtype is not None:
dtype = int(_np_dtype_to_onnx(ctx.out_dtype))
return ctx.op(op_type, [ctx.x], dtype=dtype, **attrs)
@register("randn", "size", methods=False)
def _handle_randn(ctx: OpContext) -> str:
return _drawn(ctx, "RandomNormal", mean=0.0, scale=1.0)
@register("rand", "size", methods=False)
def _handle_rand(ctx: OpContext) -> str:
return _drawn(ctx, "RandomUniform", low=0.0, high=1.0)
@register("randn_like", "input dtype device requires_grad", methods=False)
def _handle_randn_like(ctx: OpContext) -> str:
return _drawn_like(ctx, "RandomNormalLike", mean=0.0, scale=1.0)
@register("rand_like", "input dtype device requires_grad", methods=False)
def _handle_rand_like(ctx: OpContext) -> str:
return _drawn_like(ctx, "RandomUniformLike", low=0.0, high=1.0)
# ---------------------------------------------------------------------------
# Export-time shape guards
# ---------------------------------------------------------------------------
def _guard_noop(ctx: OpContext) -> None:
"""Lower an export-time shape guard to nothing.
The guards validate dynamic dimension constraints when the exported
program runs eagerly; an ONNX model declares those dimensions
symbolically instead, so the guards carry no information for export.
Their results are never consumed by other nodes.
"""
for _guard_name, _guard_params in (
("_assert_dim_range", "tensor index min max name"),
("_assert_dims_equal", "tensor_a index_a tensor_b index_b name"),
("_assert_shape_guard", "value expected guard"),
(
"_assert_dim_relation",
"tensor_root index_root tensor_derived index_derived scale offset name",
),
):
register(
_guard_name,
_guard_params,
module="tensorplay.export._trace",
methods=False,
)(_guard_noop)
# ---------------------------------------------------------------------------
# Tensor metadata
# ---------------------------------------------------------------------------
@register("getattr", "input name default", module="builtins", methods=False)
def _handle_getattr(ctx: OpContext) -> Any:
"""A tensor attribute read on a value the graph computed.
Extents come back as :func:`_dims` gives them: fixed ones as ints, ones
that vary at runtime as values read with ``Shape``. Metadata ONNX keeps
fixed in the graph (dtype, device, ...) is what the example run produced.
"""
data, attribute = ctx.x, ctx.get("name")
if attribute == "shape":
return tuple(_dims(ctx, data))
if attribute == "ndim":
return ctx.rank(data)
if attribute in ("T", "mT"):
rank = ctx.rank(data)
perm = list(range(rank))
if attribute == "T":
perm.reverse()
elif rank >= 2:
perm[-2], perm[-1] = perm[-1], perm[-2]
else:
raise UnsupportedOperatorError(
f"{ctx.node_name}: mT needs a tensor of rank 2 or more, got {rank}"
)
return ctx.op("Transpose", [data], perm=perm)
if ctx.sample is not None and not _is_tensor_like(ctx.sample):
return ctx.sample
raise UnsupportedOperatorError(
f"{ctx.node_name}: the tensor attribute {attribute!r} has no ONNX "
"lowering; read shape, ndim, T, mT or fixed metadata such as dtype"
)
@register_method("size", "input dim")
def _handle_size(ctx: OpContext) -> Any:
dim = ctx.get("dim")
if dim is None:
return tuple(_dims(ctx, ctx.x))
return _dims(ctx, ctx.x, [_normalize_axis(dim, ctx.rank(ctx.x))])[0]
@register_method("dim", "input")
@register_method("ndimension", "input")
def _handle_dim(ctx: OpContext) -> int:
return ctx.rank(ctx.x)
@register("numel", "input")
@register_method("nelement", "input")
def _handle_numel(ctx: OpContext) -> Any:
fixed = _fixed_dims(ctx, ctx.x)
if None not in fixed:
return int(np.prod(fixed, dtype=np.int64))
return Value(
ctx.op("ReduceProd", [ctx.op("Shape", [ctx.x])], keepdims=0), (), np.int64
)
# ---------------------------------------------------------------------------
# Indexing
# ---------------------------------------------------------------------------
@register("getitem", "input index", methods=False)
def _handle_getitem(ctx: OpContext) -> Any:
data = ctx.x
key = ctx.get("index")
if isinstance(data, (list, tuple)):
return data[key] if isinstance(key, slice) else data[int(key)]
keys = key if isinstance(key, tuple) else (key,)
rank = ctx.rank(data)
explicit = sum(1 for item in keys if item is not None and item is not Ellipsis)
expanded: list[Any] = []
for item in keys:
if item is Ellipsis:
expanded.extend([slice(None)] * (rank - explicit))
else:
expanded.append(item)
starts: list[int] = []
ends: list[int] = []
axes: list[int] = []
steps: list[int] = []
squeeze_axes: list[int] = []
unsqueeze_axes: list[int] = []
gathers: list[tuple[int, Any]] = []
axis = 0
output_axis = 0
for item in expanded:
if item is None:
unsqueeze_axes.append(output_axis)
output_axis += 1
continue
if isinstance(item, slice):
if item != slice(None):
start = 0 if item.start is None else int(item.start)
stop = _INT64_MAX if item.stop is None else int(item.stop)
step = 1 if item.step is None else int(item.step)
starts.append(start)
ends.append(stop)
axes.append(axis)
steps.append(step)
axis += 1
output_axis += 1
continue
if isinstance(item, (int, np.integer)) and not isinstance(item, bool):
index = int(item)
starts.append(index)
ends.append(_INT64_MAX if index == -1 else index + 1)
axes.append(axis)
steps.append(1)
squeeze_axes.append(axis)
axis += 1
continue
gathers.append((axis, item))
axis += 1
output_axis += 1
result: Any = data
if axes:
result = _slice(ctx, result, starts, ends, axes, steps)
for gather_axis, index in gathers:
result = ctx.op(
"Gather", [result, _cast(ctx, index, np.int64)], axis=gather_axis
)
if squeeze_axes:
result = _squeeze(ctx, result, sorted(squeeze_axes))
if unsqueeze_axes:
result = _unsqueeze(ctx, result, sorted(unsqueeze_axes))
return resultHelp improve this page
Found an error, an unclear step, or a missing example?
Was this page helpful?