yaourt.ortops.fused_kernel.reference_ops

Python reference implementations of the fused-kernel CUDA custom ops.

These OpRun subclasses provide CPU-only reference kernels for every operator registered in the yaourt.ortops.fused_kernel.cuda domain.

All kernels are listed in ALL_OPS and are pre-registered in default_ops, so models using fused-kernel CUDA ops can be evaluated on CPU without a GPU or the compiled shared library — no new_ops argument is required:

<<<

import numpy as np
import onnx.helper as oh
import onnx
from yaourt.reference import ExtendedReferenceEvaluator

TFLOAT = onnx.TensorProto.FLOAT
DOMAIN = "yaourt.ortops.fused_kernel.cuda"
model = oh.make_model(
    oh.make_graph(
        [oh.make_node("MulMul", ["A", "B", "C"], ["Z"], domain=DOMAIN)],
        "mulmul_graph",
        [oh.make_tensor_value_info(n, TFLOAT, [None]) for n in "ABC"],
        [oh.make_tensor_value_info("Z", TFLOAT, [None])],
    ),
    opset_imports=[oh.make_opsetid("", 18), oh.make_opsetid(DOMAIN, 1)],
    ir_version=10,
)
ref = ExtendedReferenceEvaluator(model)
a = np.array([1.0, 2.0, 3.0], dtype=np.float32)
b = np.array([4.0, 5.0, 6.0], dtype=np.float32)
c = np.array([7.0, 8.0, 9.0], dtype=np.float32)
(result,) = ref.run(None, {"A": a, "B": b, "C": c})
print(result)

>>>

    [ 28.  80. 162.]

Individual kernels can also be passed explicitly via new_ops when only a subset of operators is needed.

All classes set op_schema = None so that the ONNX reference runtime does not attempt to validate attributes against a schema that does not exist for custom-domain operators.

Operators provided

Unary (1 input → 1 output):

• NegXplus1 — 1 - x

• ReplaceZero — replace zero elements with scalar attribute by

• MulSigmoid — x * sigmoid(x) (Swish)

• Transpose2DCastFP16 — transpose 2-D float32 → float16

• Transpose2DCastFP32 — transpose 2-D float16 → float32

Binary (2 inputs → 1 output):

• MulMulSigmoid — x * y * sigmoid(y)

Ternary (3 inputs → 1 output):

• AddMul — (A + B) * C

• MulAdd — A * B + C

• SubMul — (A - B) * C

• MulSub — A * B - C

• AddAdd — A + B + C

• MulMul — A * B * C

Ternary (3 inputs → 2 outputs):

• AddSharedInput — (A + B, A + C)

• MulSharedInput — (A * B, A * C)

Quaternary (4 inputs → 1 output):

• AddAddAdd — A + B + C + D

• MulMulMul — A * B * C * D

Other:

• Rotary — rotary positional embedding (RoPE)

• ScatterNDOfShape — scatter into a zero tensor

• MaskedScatterNDOfShape — scatter with index masking

• TriMatrix — triangular matrix from scalar constants

class yaourt.ortops.fused_kernel.reference_ops.AddAdd(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes A + B + C element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor; must be broadcastable with A + B.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.AddAddAdd(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes A + B + C + D element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor; must be broadcastable with A + B.

• D – fourth input tensor; must be broadcastable with A + B + C.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.AddMul(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes (A + B) * C element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor (scale); must be broadcastable with A + B.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.AddSharedInput(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes (A + B, A + C) element-wise, producing two outputs.

Parameters:

• A – shared input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor; must be broadcastable with A.

Returns:

tuple (A + B, A + C), each with the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.MaskedScatterNDOfShape(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Scatters updates into a zero tensor, skipping masked index entries.

Index entries equal to maskedValue are ignored, leaving the corresponding output positions at zero.

Parameters:

• shape – 1-D int64 tensor defining the output shape.

• indices – integer indices tensor; the last dimension gives the index depth into the output tensor.

• updates – data tensor to scatter into the output.

• reduction – string attribute controlling conflict resolution; one of "add" (default), "none", "mul", "min", "max".

• maskedValue – integer attribute; index entries equal to this value are skipped (default -1).

Returns:

output tensor of dtype matching updates and shape given by shape, filled with scattered values.

class yaourt.ortops.fused_kernel.reference_ops.MulAdd(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes A * B + C element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – bias tensor; must be broadcastable with A * B.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.MulMul(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes A * B * C element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor; must be broadcastable with A * B.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.MulMulMul(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes A * B * C * D element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor; must be broadcastable with A * B.

• D – fourth input tensor; must be broadcastable with A * B * C.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.MulMulSigmoid(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes x * y * sigmoid(y) element-wise.

Parameters:

• X – first input tensor.

• Y – second input tensor (used for both the multiplication and the sigmoid gate); must be broadcastable with X.

Returns:

output tensor of the same shape and dtype as X.

class yaourt.ortops.fused_kernel.reference_ops.MulSharedInput(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes (A * B, A * C) element-wise, producing two outputs.

Parameters:

• A – shared input tensor.

• B – second input tensor; must be broadcastable with A.

• C – third input tensor; must be broadcastable with A.

Returns:

tuple (A * B, A * C), each with the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.MulSigmoid(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes x * sigmoid(x) element-wise (Swish / SiLU activation).

Parameters:

X – input tensor (float32 or float64).

Returns:

output tensor of the same shape and dtype as X.

class yaourt.ortops.fused_kernel.reference_ops.MulSub(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes A * B - C element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – bias tensor to subtract; must be broadcastable with A * B.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.NegXplus1(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes 1 - x element-wise.

Parameters:

X – input tensor (any numeric dtype).

Returns:

output tensor of the same shape and dtype as X.

class yaourt.ortops.fused_kernel.reference_ops.ReplaceZero(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Replaces every zero element with the scalar attribute by.

Parameters:

• X – input tensor (any numeric dtype).

• by – scalar replacement value for zero elements (default 0.0).

Returns:

output tensor of the same shape and dtype as X with zeros replaced by by.

class yaourt.ortops.fused_kernel.reference_ops.Rotary(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Applies a rotary positional transformation to the last dimension of X.

The side attribute controls which half of the rotation is computed:

• left: out[..., :half] = x[..., half:], out[..., half:] = -x[..., :half]

• right: out[..., :half] = -x[..., half:], out[..., half:] = x[..., :half]

Parameters:

• X – input tensor; the last dimension must be 2 * half.

• splits – 1-D int64 tensor [half, half]; the first element provides the half-size of the last dimension.

• side – string attribute "left" (default) or "right" selecting the rotation direction.

Returns:

output tensor of the same shape and dtype as X.

class yaourt.ortops.fused_kernel.reference_ops.ScatterNDOfShape(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Scatters updates into a zero tensor of shape shape.

Parameters:

• shape – 1-D int64 tensor defining the output shape.

• indices – integer indices tensor; the last dimension gives the index depth into the output tensor.

• updates – data tensor to scatter into the output.

• reduction – string attribute controlling conflict resolution; one of "add" (default), "none", "mul", "min", "max".

Returns:

output tensor of dtype matching updates and shape given by shape, filled with scattered values.

class yaourt.ortops.fused_kernel.reference_ops.SubMul(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Computes (A - B) * C element-wise.

Parameters:

• A – first input tensor.

• B – second input tensor; must be broadcastable with A.

• C – scale tensor; must be broadcastable with A - B.

Returns:

output tensor of the same shape and dtype as A.

class yaourt.ortops.fused_kernel.reference_ops.Transpose2DCastFP16(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Transposes a 2-D float32 matrix and casts the result to float16.

Parameters:

X – 2-D input tensor of dtype float32 with shape (M, N).

Returns:

2-D output tensor of dtype float16 with shape (N, M).

class yaourt.ortops.fused_kernel.reference_ops.Transpose2DCastFP32(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Transposes a 2-D float16 matrix and casts the result to float32.

Parameters:

X – 2-D input tensor of dtype float16 with shape (M, N).

Returns:

2-D output tensor of dtype float32 with shape (N, M).

class yaourt.ortops.fused_kernel.reference_ops.TriMatrix(onnx_node: NodeProto, run_params: dict[str, Any], schema: Any = None)

Fills a 2-D matrix whose elements depend on their position relative to the main diagonal.

The output element at (r, c) equals:

• csts[0] when r > c (lower triangle),

• csts[1] when r == c (diagonal),

• csts[2] when r < c (upper triangle).

Parameters:

• shape – 1-D int64 tensor [n_rows, n_cols] giving the output dimensions.

• csts – 1-D tensor with exactly three values [lower_value, diag_value, upper_value].

Returns:

2-D output tensor with shape (n_rows, n_cols) and dtype matching csts.

Usage example

<<<

import numpy as np
import onnx.helper as oh
import onnx
from yaourt.reference import ExtendedReferenceEvaluator

TFLOAT = onnx.TensorProto.FLOAT
DOMAIN = "yaourt.ortops.fused_kernel.cuda"
model = oh.make_model(
    oh.make_graph(
        [oh.make_node("MulMul", ["A", "B", "C"], ["Z"], domain=DOMAIN)],
        "mulmul_graph",
        [oh.make_tensor_value_info(n, TFLOAT, [None]) for n in "ABC"],
        [oh.make_tensor_value_info("Z", TFLOAT, [None])],
    ),
    opset_imports=[oh.make_opsetid("", 18), oh.make_opsetid(DOMAIN, 1)],
    ir_version=10,
)
ref = ExtendedReferenceEvaluator(model)
a = np.array([1.0, 2.0, 3.0], dtype=np.float32)
b = np.array([4.0, 5.0, 6.0], dtype=np.float32)
c = np.array([7.0, 8.0, 9.0], dtype=np.float32)
(result,) = ref.run(None, {"A": a, "B": b, "C": c})
print(result)

>>>

    [ 28.  80. 162.]