onnx--onnx/docs/TestCoverage.md

Test Coverage Report (ONNX Core Operators)

Outlines

• Node Test Coverage
• Model Test Coverage
• Overall Test Coverage

Node Test Coverage

Summary

Node tests have covered 189/201 (94.03%, 5 generators excluded) common operators.

Node tests have covered 1/1 (100.00%, 0 generators excluded) experimental operators.

• Covered Common Operators
• No Cover Common Operators
• Covered Experimental Operators
• No Cover Experimental Operators

💚Covered Common Operators

Abs

There are 1 test cases, listed as following:

abs

node = onnx.helper.make_node(
    "Abs",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.abs(x)

expect(node, inputs=[x], outputs=[y], name="test_abs")

Acos

There are 1 test cases, listed as following:

acos

node = onnx.helper.make_node(
    "Acos",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-0.5, 0, 0.5]).astype(np.float32)
y = np.arccos(x)
expect(node, inputs=[x], outputs=[y], name="test_acos_example")

x = np.random.rand(3, 4, 5).astype(np.float32)
y = np.arccos(x)
expect(node, inputs=[x], outputs=[y], name="test_acos")

Acosh

There are 1 test cases, listed as following:

acosh

node = onnx.helper.make_node(
    "Acosh",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([10, np.e, 1]).astype(np.float32)
y = np.arccosh(x)  # expected output [2.99322295,  1.65745449,  0.]
expect(node, inputs=[x], outputs=[y], name="test_acosh_example")

x = np.random.uniform(1.0, 10.0, (3, 4, 5)).astype(np.float32)
y = np.arccosh(x)
expect(node, inputs=[x], outputs=[y], name="test_acosh")

Adagrad

There are 2 test cases, listed as following:

adagrad

# Define operator attributes.
norm_coefficient = 0.001
epsilon = 1e-5
decay_factor = 0.1

# Create operator.
node = onnx.helper.make_node(
    "Adagrad",
    inputs=["R", "T", "X", "G", "H"],
    outputs=["X_new", "H_new"],
    norm_coefficient=norm_coefficient,
    epsilon=epsilon,
    decay_factor=decay_factor,
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar
x = np.array([1.0], dtype=np.float32)
g = np.array([-1.0], dtype=np.float32)
h = np.array([2.0], dtype=np.float32)

# Compute expected outputs of Adagrad.
x_new, h_new = apply_adagrad(
    r, t, x, g, h, norm_coefficient, epsilon, decay_factor
)

# Check results.
expect(
    node,
    inputs=[r, t, x, g, h],
    outputs=[x_new, h_new],
    name="test_adagrad",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

adagrad_multiple

# Define operator attributes.
norm_coefficient = 0.001
epsilon = 1e-5
decay_factor = 0.1

node = onnx.helper.make_node(
    "Adagrad",
    inputs=["R", "T", "X1", "X2", "G1", "G2", "H1", "H2"],
    outputs=["X1_new", "X2_new", "H1_new", "H2_new"],
    norm_coefficient=norm_coefficient,
    epsilon=epsilon,
    decay_factor=decay_factor,
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar

x1 = np.array([1.0], dtype=np.float32)
g1 = np.array([-1.0], dtype=np.float32)
h1 = np.array([2.0], dtype=np.float32)

x2 = np.array([1.0, 2.0], dtype=np.float32)
g2 = np.array([-1.0, -3.0], dtype=np.float32)
h2 = np.array([4.0, 1.0], dtype=np.float32)

# Compute expected outputs of Adagrad.
x1_new, h1_new = apply_adagrad(
    r, t, x1, g1, h1, norm_coefficient, epsilon, decay_factor
)
x2_new, h2_new = apply_adagrad(
    r, t, x2, g2, h2, norm_coefficient, epsilon, decay_factor
)

# Check results.
expect(
    node,
    inputs=[r, t, x1, x2, g1, g2, h1, h2],
    outputs=[x1_new, x2_new, h1_new, h2_new],
    name="test_adagrad_multiple",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

Adam

There are 2 test cases, listed as following:

adam

# Define operator attributes.
norm_coefficient = 0.001
alpha = 0.95
beta = 0.1
epsilon = 1e-7

# Create operator.
node = onnx.helper.make_node(
    "Adam",
    inputs=["R", "T", "X", "G", "V", "H"],
    outputs=["X_new", "V_new", "H_new"],
    norm_coefficient=norm_coefficient,
    alpha=alpha,
    beta=beta,
    epsilon=epsilon,
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar
x = np.array([1.2, 2.8], dtype=np.float32)
g = np.array([-0.94, -2.5], dtype=np.float32)
v = np.array([1.7, 3.6], dtype=np.float32)
h = np.array([0.1, 0.1], dtype=np.float32)

# Compute expected outputs of Adam.
x_new, v_new, h_new = apply_adam(
    r, t, x, g, v, h, norm_coefficient, 0.0, alpha, beta, epsilon
)

# Check results.
expect(
    node,
    inputs=[r, t, x, g, v, h],
    outputs=[x_new, v_new, h_new],
    name="test_adam",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

adam_multiple

# Define operator attributes.
norm_coefficient = 0.001
alpha = 0.95
beta = 0.85
epsilon = 1e-2

node = onnx.helper.make_node(
    "Adam",
    inputs=["R", "T", "X1", "X2", "G1", "G2", "V1", "V2", "H1", "H2"],
    outputs=["X1_new", "X2_new", "V1_new", "V2_new", "H1_new", "H2_new"],
    norm_coefficient=norm_coefficient,
    alpha=alpha,
    beta=beta,
    epsilon=epsilon,
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar

x1 = np.array([1.0], dtype=np.float32)
g1 = np.array([-1.0], dtype=np.float32)
v1 = np.array([2.0], dtype=np.float32)
h1 = np.array([0.5], dtype=np.float32)

x2 = np.array([1.0, 2.0], dtype=np.float32)
g2 = np.array([-1.0, -3.0], dtype=np.float32)
v2 = np.array([4.0, 1.0], dtype=np.float32)
h2 = np.array([1.0, 10.0], dtype=np.float32)

# Compute expected outputs of Adam.
x1_new, v1_new, h1_new = apply_adam(
    r, t, x1, g1, v1, h1, norm_coefficient, 0.0, alpha, beta, epsilon
)
x2_new, v2_new, h2_new = apply_adam(
    r, t, x2, g2, v2, h2, norm_coefficient, 0.0, alpha, beta, epsilon
)

# Check results.
expect(
    node,
    inputs=[r, t, x1, x2, g1, g2, v1, v2, h1, h2],
    outputs=[x1_new, x2_new, v1_new, v2_new, h1_new, h2_new],
    name="test_adam_multiple",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

Add

There are 2 test cases, listed as following:

add

node = onnx.helper.make_node(
    "Add",
    inputs=["x", "y"],
    outputs=["sum"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_int8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_int16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint64")

add_broadcast

node = onnx.helper.make_node(
    "Add",
    inputs=["x", "y"],
    outputs=["sum"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
expect(node, inputs=[x, y], outputs=[x + y], name="test_add_bcast")

AffineGrid

There are 2 test cases, listed as following:

2d_no_reference_evaluator

theta_2d = create_theta_2d()
N, C, H, W = len(theta_2d), 3, 5, 6
data_size = (H, W)
for align_corners in (0, 1):
    node = onnx.helper.make_node(
        "AffineGrid",
        inputs=["theta", "size"],
        outputs=["grid"],
        align_corners=align_corners,
    )

    original_grid = construct_original_grid(data_size, align_corners)
    grid = apply_affine_transform(theta_2d, original_grid)

    test_name = "test_affine_grid_2d"
    if align_corners == 1:
        test_name += "_align_corners"
    expect(
        node,
        inputs=[theta_2d, np.array([N, C, H, W], dtype=np.int64)],
        outputs=[grid],
        name=test_name,
    )

3d_no_reference_evaluator

theta_3d = create_theta_3d()
N, C, D, H, W = len(theta_3d), 3, 4, 5, 6
data_size = (D, H, W)
for align_corners in (0, 1):
    node = onnx.helper.make_node(
        "AffineGrid",
        inputs=["theta", "size"],
        outputs=["grid"],
        align_corners=align_corners,
    )

    original_grid = construct_original_grid(data_size, align_corners)
    grid = apply_affine_transform(theta_3d, original_grid)

    test_name = "test_affine_grid_3d"
    if align_corners == 1:
        test_name += "_align_corners"
    expect(
        node,
        inputs=[theta_3d, np.array([N, C, D, H, W], dtype=np.int64)],
        outputs=[grid],
        name=test_name,
    )

And

There are 2 test cases, listed as following:

and

node = onnx.helper.make_node(
    "And",
    inputs=["x", "y"],
    outputs=["and"],
)

# 2d
x = (np.random.randn(3, 4) > 0).astype(bool)
y = (np.random.randn(3, 4) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and2d")

# 3d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(3, 4, 5) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and3d")

# 4d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and4d")

and_broadcast

node = onnx.helper.make_node(
    "And",
    inputs=["x", "y"],
    outputs=["and"],
)

# 3d vs 1d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(5) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and_bcast3v1d")

# 3d vs 2d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(4, 5) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and_bcast3v2d")

# 4d vs 2d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(5, 6) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and_bcast4v2d")

# 4d vs 3d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(4, 5, 6) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and_bcast4v3d")

# 4d vs 4d
x = (np.random.randn(1, 4, 1, 6) > 0).astype(bool)
y = (np.random.randn(3, 1, 5, 6) > 0).astype(bool)
z = np.logical_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_and_bcast4v4d")

ArgMax

There are 8 test cases, listed as following:

default_axes_keepdims

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
keepdims = 1
node = onnx.helper.make_node(
    "ArgMax", inputs=["data"], outputs=["result"], keepdims=keepdims
)

# result: [[1, 1]]
result = argmax_use_numpy(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_default_axis_example",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [1, 3, 4]
result = argmax_use_numpy(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_default_axis_random",
)

default_axes_keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
keepdims = 1
node = onnx.helper.make_node(
    "ArgMax",
    inputs=["data"],
    outputs=["result"],
    keepdims=keepdims,
    select_last_index=True,
)

# result: [[1, 1]]
result = argmax_use_numpy_select_last_index(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_default_axis_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [1, 3, 4]
result = argmax_use_numpy_select_last_index(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_default_axis_random_select_last_index",
)

keepdims

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMax", inputs=["data"], outputs=["result"], axis=axis, keepdims=keepdims
)
# result: [[0], [1]]
result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node, inputs=[data], outputs=[result], name="test_argmax_keepdims_example"
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 1, 4]
result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node, inputs=[data], outputs=[result], name="test_argmax_keepdims_random"
)

keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMax",
    inputs=["data"],
    outputs=["result"],
    axis=axis,
    keepdims=keepdims,
    select_last_index=True,
)
# result: [[1], [1]]
result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_keepdims_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 1, 4]
result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_keepdims_random_select_last_index",
)

negative_axis_keepdims

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = -1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMax", inputs=["data"], outputs=["result"], axis=axis, keepdims=keepdims
)
# result: [[0], [1]]
result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_negative_axis_keepdims_example",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 3, 1]
result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_negative_axis_keepdims_random",
)

negative_axis_keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = -1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMax",
    inputs=["data"],
    outputs=["result"],
    axis=axis,
    keepdims=keepdims,
    select_last_index=True,
)
# result: [[1], [1]]
result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_negative_axis_keepdims_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 3, 1]
result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_negative_axis_keepdims_random_select_last_index",
)

no_keepdims

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 0
node = onnx.helper.make_node(
    "ArgMax", inputs=["data"], outputs=["result"], axis=axis, keepdims=keepdims
)
# result: [0, 1]
result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_no_keepdims_example",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 4]
result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node, inputs=[data], outputs=[result], name="test_argmax_no_keepdims_random"
)

no_keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 0
node = onnx.helper.make_node(
    "ArgMax",
    inputs=["data"],
    outputs=["result"],
    axis=axis,
    keepdims=keepdims,
    select_last_index=True,
)
# result: [1, 1]
result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_no_keepdims_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 4]
result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmax_no_keepdims_random_select_last_index",
)

ArgMin

There are 8 test cases, listed as following:

default_axes_keepdims

data = np.array([[2, 1], [3, 10]], dtype=np.float32)
keepdims = 1
node = onnx.helper.make_node(
    "ArgMin", inputs=["data"], outputs=["result"], keepdims=keepdims
)

# The content of result is : [[0], [0]]
result = argmin_use_numpy(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_default_axis_example",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [1, 3, 4]
result = argmin_use_numpy(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_default_axis_random",
)

default_axes_keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
keepdims = 1
node = onnx.helper.make_node(
    "ArgMin",
    inputs=["data"],
    outputs=["result"],
    keepdims=keepdims,
    select_last_index=True,
)

# result: [[0, 0]]
result = argmin_use_numpy_select_last_index(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_default_axis_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [1, 3, 4]
result = argmin_use_numpy_select_last_index(data, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_default_axis_random_select_last_index",
)

keepdims

data = np.array([[2, 1], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMin", inputs=["data"], outputs=["result"], axis=axis, keepdims=keepdims
)
# The content of result is : [[1], [0]]
result = argmin_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node, inputs=[data], outputs=[result], name="test_argmin_keepdims_example"
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 1, 4]
result = argmin_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node, inputs=[data], outputs=[result], name="test_argmin_keepdims_random"
)

keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMin",
    inputs=["data"],
    outputs=["result"],
    axis=axis,
    keepdims=keepdims,
    select_last_index=True,
)
# result: [[1], [0]]
result = argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_keepdims_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 1, 4]
result = argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_keepdims_random_select_last_index",
)

negative_axis_keepdims

data = np.array([[2, 1], [3, 10]], dtype=np.float32)
axis = -1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMin", inputs=["data"], outputs=["result"], axis=axis, keepdims=keepdims
)
# The content of result is : [[1], [0]]
result = argmin_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_negative_axis_keepdims_example",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 3, 1]
result = argmin_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_negative_axis_keepdims_random",
)

negative_axis_keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = -1
keepdims = 1
node = onnx.helper.make_node(
    "ArgMin",
    inputs=["data"],
    outputs=["result"],
    axis=axis,
    keepdims=keepdims,
    select_last_index=True,
)
# result: [[1], [0]]
result = argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_negative_axis_keepdims_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 3, 1]
result = argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_negative_axis_keepdims_random_select_last_index",
)

no_keepdims

data = np.array([[2, 1], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 0
node = onnx.helper.make_node(
    "ArgMin", inputs=["data"], outputs=["result"], axis=axis, keepdims=keepdims
)
# The content of result is : [[1, 0]]
result = argmin_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_no_keepdims_example",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 4]
result = argmin_use_numpy(data, axis=axis, keepdims=keepdims)
expect(
    node, inputs=[data], outputs=[result], name="test_argmin_no_keepdims_random"
)

no_keepdims_select_last_index

data = np.array([[2, 2], [3, 10]], dtype=np.float32)
axis = 1
keepdims = 0
node = onnx.helper.make_node(
    "ArgMin",
    inputs=["data"],
    outputs=["result"],
    axis=axis,
    keepdims=keepdims,
    select_last_index=True,
)
# result: [[1, 0]]
result = argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_no_keepdims_example_select_last_index",
)

data = np.random.uniform(-10, 10, [2, 3, 4]).astype(np.float32)
# result's shape: [2, 4]
result = argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
expect(
    node,
    inputs=[data],
    outputs=[result],
    name="test_argmin_no_keepdims_random_select_last_index",
)

Asin

There are 1 test cases, listed as following:

asin

node = onnx.helper.make_node(
    "Asin",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-0.5, 0, 0.5]).astype(np.float32)
y = np.arcsin(x)
expect(node, inputs=[x], outputs=[y], name="test_asin_example")

x = np.random.rand(3, 4, 5).astype(np.float32)
y = np.arcsin(x)
expect(node, inputs=[x], outputs=[y], name="test_asin")

Asinh

There are 1 test cases, listed as following:

asinh

node = onnx.helper.make_node(
    "Asinh",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.arcsinh(x)  # expected output [-0.88137358,  0.,  0.88137358]
expect(node, inputs=[x], outputs=[y], name="test_asinh_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.arcsinh(x)
expect(node, inputs=[x], outputs=[y], name="test_asinh")

Atan

There are 1 test cases, listed as following:

atan

node = onnx.helper.make_node(
    "Atan",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.arctan(x)
expect(node, inputs=[x], outputs=[y], name="test_atan_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.arctan(x)
expect(node, inputs=[x], outputs=[y], name="test_atan")

Atanh

There are 1 test cases, listed as following:

atanh

node = onnx.helper.make_node(
    "Atanh",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-0.5, 0, 0.5]).astype(np.float32)
y = np.arctanh(x)  # expected output [-0.54930615,  0.,  0.54930615]
expect(node, inputs=[x], outputs=[y], name="test_atanh_example")

x = np.random.uniform(0.0, 1.0, (3, 4, 5)).astype(np.float32)
y = np.arctanh(x)
expect(node, inputs=[x], outputs=[y], name="test_atanh")

Attention

There are 76 test cases, listed as following:

attention

node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"])

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_23_boolmask_fullymasked_row_nan_robustness

"""Opset-23 fully-masked boolean ``attn_mask`` row -> zero (not ``NaN``).

This locks the opset-23 / ``old.cc`` function-body fully-masked-row guard
against future regressions. In opset 23 the only in-contract fully-masked
row comes from an all-``False`` boolean ``attn_mask`` row (``is_causal`` is
not set here): every key for that query is disallowed, so ``softmax`` over an
all-``-inf`` bias row is ``NaN``. The guard zeros that row's probabilities
before the ``P @ V`` contraction so the output row is exactly ``0``, while
rows with at least one allowed key are unchanged.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(4)
B, H, S, D = 1, 2, 2, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(B, H, S, D).astype(np.float32)
K = np.random.rand(B, H, S, D).astype(np.float32)
V = np.random.rand(B, H, S, D).astype(np.float32)
# Row 0: no key allowed -> fully masked (Bug-2 empty row). Row 1: both keys
# allowed -> finite, unchanged by the guard.
attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_)

Y, _, _, _ = _compute_attention(Q, K, V, attn_mask=attn_mask)

# Fully-masked row 0 is exactly zero (not NaN); every other cell is finite.
assert np.all(np.isfinite(Y)), "non-masked rows must be finite"
assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), (
    "fully-masked row must be zero (Bug-2)"
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_23_boolmask_fullymasked_row_nan_robustness",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_23_fullymasked_qk_matmul_output_mode3_zero

"""Opset-23 ``qk_matmul_output_mode=3`` fully-masked row is a zero row.

Mode ``3`` exposes the post-softmax matrix as the optional
``qk_matmul_output``.  For a fully-masked query row (all-``False`` boolean
``attn_mask`` row), the fully-masked-row guard is applied before this output
is produced, so the mode-3 row is zeroed, consistent with the primary output
``Y`` row (both are ``0``).  This pins the mandated agreement between the
guarded primary output and the mode-3 output at opset 23.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(13)
B, H, S, D = 1, 2, 2, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y", "", "", "qk_matmul_output"],
    qk_matmul_output_mode=3,
)

Q = np.random.rand(B, H, S, D).astype(np.float32)
K = np.random.rand(B, H, S, D).astype(np.float32)
V = np.random.rand(B, H, S, D).astype(np.float32)
# Row 0: no key allowed -> fully masked. Row 1: both keys allowed -> finite.
attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_)

Y, _, _, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    qk_matmul_output_mode=3,
)

# Primary output row 0 and the mode-3 row 0 are both guarded to zero.
assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), (
    "fully-masked primary output row must be zero"
)
assert np.array_equal(
    qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :])
), "mode-3 output row for a fully-masked query must be zero (consistent with Y)"
assert np.all(np.isfinite(qk_matmul_output)), (
    "all mode-3 rows are finite (the fully-masked row is guarded to 0.0)"
)
assert np.all(np.isfinite(Y)), (
    "all Y rows are finite (the fully-masked row is guarded to 0.0)"
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y, qk_matmul_output],
    name="test_attention_23_fullymasked_qk_matmul_output_mode3_zero",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_24_fullymasked_qk_matmul_output_mode3_zero

"""Opset-24 ``qk_matmul_output_mode=3`` fully-masked row is a zero row.

The opset-24 twin of
``export_attention_23_fullymasked_qk_matmul_output_mode3_zero``.  Mode ``3``
exposes the post-softmax matrix as the optional ``qk_matmul_output``.  For a
fully-masked query row (all-``False`` boolean ``attn_mask`` row), the
fully-masked-row guard is applied before this output is produced, so the
mode-3 row is zeroed, consistent with the primary output ``Y`` row (both are
``0``).  This pins the mandated agreement between the guarded primary output
and the mode-3 output at opset 24.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(13)
B, H, S, D = 1, 2, 2, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y", "", "", "qk_matmul_output"],
    qk_matmul_output_mode=3,
)

Q = np.random.rand(B, H, S, D).astype(np.float32)
K = np.random.rand(B, H, S, D).astype(np.float32)
V = np.random.rand(B, H, S, D).astype(np.float32)
# Row 0: no key allowed -> fully masked. Row 1: both keys allowed -> finite.
attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_)

Y, _, _, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    qk_matmul_output_mode=3,
)

# Primary output row 0 and the mode-3 row 0 are both guarded to zero.
assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), (
    "fully-masked primary output row must be zero"
)
assert np.array_equal(
    qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :])
), "mode-3 output row for a fully-masked query must be zero (consistent with Y)"
assert np.all(np.isfinite(qk_matmul_output)), (
    "all mode-3 rows are finite (the fully-masked row is guarded to 0.0)"
)
assert np.all(np.isfinite(Y)), (
    "all Y rows are finite (the fully-masked row is guarded to 0.0)"
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y, qk_matmul_output],
    name="test_attention_24_fullymasked_qk_matmul_output_mode3_zero",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_24_qk_matmul_output_mode3_softmax_precision

"""Mode-3 ``qk_matmul_output`` is emitted at the output precision ``T1``.

``qk_matmul_output_mode=3`` exposes the post-softmax probabilities.  When
``softmax_precision`` differs from the operator's output type ``T1`` (here
``T1 = float16`` with softmax computed in ``float32``), the mode-3 output is
cast back to ``T1`` -- matching the reference implementation, which casts the
exposed matrix to ``Q.dtype``.  This locks both the dtype contract and the
fully-masked-row zeroing under a non-default ``softmax_precision``.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(17)
B, H, S, D = 1, 2, 2, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y", "", "", "qk_matmul_output"],
    qk_matmul_output_mode=3,
    softmax_precision=int(onnx.TensorProto.FLOAT),
)

# T1 = float16; softmax runs in float32, so the mode-3 output is cast back to
# float16 on emission.
Q = np.random.rand(B, H, S, D).astype(np.float16)
K = np.random.rand(B, H, S, D).astype(np.float16)
V = np.random.rand(B, H, S, D).astype(np.float16)
# Row 0: fully masked. Row 1: both keys allowed -> finite.
attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_)

Y, _, _, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    qk_matmul_output_mode=3,
    softmax_precision=int(onnx.TensorProto.FLOAT),
)

# The mode-3 output is emitted at T1 (float16), not the float32 softmax
# precision, matching the operator's output type.
assert qk_matmul_output.dtype == np.float16, (
    "mode-3 qk_matmul_output must be emitted at the output precision T1 (float16)"
)
# The fully-masked row is still guarded to zero, consistent with Y.
assert np.array_equal(
    qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :])
), "mode-3 output row for a fully-masked query must be zero (consistent with Y)"
assert np.all(np.isfinite(qk_matmul_output)), (
    "all mode-3 rows are finite (the fully-masked row is guarded to 0.0)"
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y, qk_matmul_output],
    name="test_attention_24_qk_matmul_output_mode3_softmax_precision",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_3d

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_attn_mask

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_3d_attn_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_causal

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    is_causal=1,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    is_causal=1,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_diff_head_sizes

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 30).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_diff_heads_sizes",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_diff_head_sizes_attn_mask

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 30).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_3d_diff_heads_sizes_attn_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_diff_head_sizes_causal

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    is_causal=1,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 30).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    is_causal=1,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_diff_heads_sizes_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_diff_head_sizes_scaled

scale = 1e-2
q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 30).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    scale=scale,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_diff_heads_sizes_scaled",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_diff_head_sizes_softcap

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    softcap=3.0,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 30).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    softcap=3.0,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_diff_heads_sizes_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_diff_head_sizes_with_past_and_present

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 30).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_3d_diff_heads_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_gqa

q_num_heads, kv_num_heads = 9, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 72).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_gqa",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_gqa_attn_mask

q_num_heads, kv_num_heads = 9, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 72).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_3d_gqa_attn_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_gqa_causal

q_num_heads, kv_num_heads = 9, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    is_causal=1,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 72).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    is_causal=1,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_gqa_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_gqa_scaled

scale = 1e-2
q_num_heads, kv_num_heads = 9, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 72).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    scale=scale,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_gqa_scaled",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_gqa_softcap

q_num_heads, kv_num_heads = 9, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    softcap=3.0,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 72).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    softcap=3.0,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_gqa_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_gqa_with_past_and_present

q_num_heads, kv_num_heads = 9, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 72).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_3d_gqa_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_scaled

scale = 1e-2
q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    scale=scale,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_scaled",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_softcap

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    softcap=3.0,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    softcap=3.0,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_transpose_verification

"""Test case to verify correct 3D to 4D transpose behavior.

This test verifies that 3D inputs are correctly reshaped and transposed
according to the ONNX specification:
[batch_size, seq_length, hidden_size] ->
[batch_size, seq_length, num_heads, head_size] ->
[batch_size, num_heads, seq_length, head_size]
"""
q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

# Test inputs that will clearly demonstrate the transpose behavior
batch_size = 1
q_seq_length = 2
kv_seq_length = 2
head_size = 4
q_hidden_size = q_num_heads * head_size  # 3 * 4 = 12
kv_hidden_size = kv_num_heads * head_size  # 3 * 4 = 12

# Create structured inputs to verify correct transpose behavior
# Q has a pattern where each position in hidden dimension has a specific value
Q = np.zeros((batch_size, q_seq_length, q_hidden_size), dtype=np.float32)
# Fill Q with pattern: head0=[1,1,1,1], head1=[2,2,2,2], head2=[3,3,3,3]
for head in range(q_num_heads):
    start_idx = head * head_size
    end_idx = start_idx + head_size
    Q[0, :, start_idx:end_idx] = float(head + 1)

K = np.ones((batch_size, kv_seq_length, kv_hidden_size), dtype=np.float32) * 0.1
V = np.ones((batch_size, kv_seq_length, kv_hidden_size), dtype=np.float32) * 0.1

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_3d_transpose_verification",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_with_past_and_present

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_3d_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_with_past_and_present_qk_matmul

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_3d_with_past_and_present_qk_matmul",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_with_past_and_present_qk_matmul_bias

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
    qk_matmul_output_mode=2,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
    qk_matmul_output_mode=2,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_3d_with_past_and_present_qk_matmul_bias",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_with_past_and_present_qk_matmul_softcap

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
    softcap=2.0,
    qk_matmul_output_mode=1,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
    softcap=2.0,
    qk_matmul_output_mode=1,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_3d_with_past_and_present_qk_matmul_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_3d_with_past_and_present_qk_matmul_softmax

q_num_heads, kv_num_heads = 3, 3
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
    qk_matmul_output_mode=3,
)

past_sequence_length = 12
Q = np.random.rand(2, 4, 24).astype(np.float32)
K = np.random.rand(2, 6, 24).astype(np.float32)
V = np.random.rand(2, 6, 24).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    q_num_heads=q_num_heads,
    kv_num_heads=kv_num_heads,
    qk_matmul_output_mode=3,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_3d_with_past_and_present_qk_matmul_softmax",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_4d_causal_nonpad_attn_mask_composition

"""Compose ``is_causal`` + ``nonpad_kv_seqlen`` + boolean ``attn_mask``.

The existing nonpad tests use no ``attn_mask`` and the existing mask tests
use no ``nonpad_kv_seqlen``; this is the first to activate all three
constraints together on the external-cache path with ``batch > 1``.  The
three biases are summed additively and a key is attended only if allowed by
all three.  Crucially the inputs are designed so that **each constraint is
independently necessary** -- removing any one changes the golden -- to avoid a
degenerate test that a backend ignoring ``is_causal`` and/or
``nonpad_kv_seqlen`` could still pass:

* **``is_causal`` binds.**  Each batch has a key that the boolean mask allows
  (``True``) but the bottom-right causal frontier disallows
  (``j > i + offset``); only ``is_causal`` masks it (batch 0 row 0 key 2,
  batch 1 row 0 key 3).
* **``attn_mask`` binds.**  Each batch has a key the causal frontier and the
  padding bound both allow but the boolean mask sets ``False`` (batch 0 row 2
  key 1, batch 1 row 2 key 2); only the mask masks it.
* **``nonpad_kv_seqlen`` binds.**  ``nonpad_kv_seqlen`` sets the per-batch
  causal *offset* (``offset = nonpad_kv_seqlen - q_sequence_length``), so
  dropping it collapses the frontier to top-left (``offset = 0``) and shifts
  which keys are attended.  (Under ``is_causal=1`` the causal frontier already
  subsumes the ``j < nonpad`` padding bound, so ``nonpad_kv_seqlen`` binds
  through the offset it induces rather than through a redundant padding cut.)

The mask is chosen to leave at least one allowed key on every query row, so
this exercises the *intersection* of the three constraints with finite outputs
(the fully-masked-row guard is covered by
``test_attention_4d_causal_nonpad_negative_offset_structural_empty`` and
``test_attention_24_fullymasked_qk_matmul_output_mode3_zero``).

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(11)
B, H, L, D = 2, 2, 6, 8
S_q = 3

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H, S_q, D).astype(np.float32)
K = np.random.rand(B, H, L, D).astype(np.float32)
V = np.random.rand(B, H, L, D).astype(np.float32)
nonpad_kv_seqlen = np.array([4, 5], dtype=np.int64)  # offsets [1, 2]
# Per-batch (B, 1, S_q, L) bool mask. Each batch is laid out so all three
# constraints uniquely bind (see the docstring): a causal-only-masked key
# (mask True, j > i + offset), a mask-only-masked key (mask False, causal +
# nonpad allow it), and >=1 allowed key per row.
attn_mask = np.array(
    [
        [
            [
                [True, True, True, False, False, False],
                [True, True, True, False, False, False],
                [True, False, True, True, False, False],
            ]
        ],
        [
            [
                [True, True, True, True, False, False],
                [True, True, True, True, False, False],
                [True, True, False, True, True, False],
            ]
        ],
    ],
    dtype=np.bool_,
)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
    is_causal=1,
)

# The chosen mask leaves >=1 allowed key per row, so the composition stays
# finite (no fully-masked row in this case).
assert np.all(np.isfinite(Y)), "composed-constraint output must be finite"

# Non-degeneracy: each constraint is independently necessary. Removing any one
# of the three (is_causal, attn_mask, nonpad_kv_seqlen) must change the result,
# so a backend that ignores is_causal or nonpad_kv_seqlen cannot reproduce the
# golden by applying only the most restrictive mask.
y_no_causal, _, _, _ = _compute_attention(
    Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=0
)
y_no_mask, _, _, _ = _compute_attention(
    Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1
)
y_no_nonpad, _, _, _ = _compute_attention(
    Q, K, V, attn_mask=attn_mask, is_causal=1
)
assert not np.allclose(Y, y_no_causal, equal_nan=True), (
    "is_causal must bind: dropping it changes the result"
)
assert not np.allclose(Y, y_no_mask, equal_nan=True), (
    "attn_mask must bind: dropping it changes the result"
)
assert not np.allclose(Y, y_no_nonpad, equal_nan=True), (
    "nonpad_kv_seqlen must bind (via the causal offset): dropping it changes the result"
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_causal_nonpad_attn_mask_composition",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_causal_nonpad_batch_prefill

"""Batch>1 continued prefill with distinct per-batch bottom-right offsets.

The batched generalization of the ``batch == 1`` continued-prefill case: with
``nonpad_kv_seqlen = [4, 5, 6]`` and ``S_q = 2`` the per-batch bottom-right
offsets are ``[2, 3, 4]`` (all ``>= 0``), so each batch realigns its causal
frontier to its own valid-key prefix.  This pins that the per-batch offset is
applied independently across the batch dimension.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(12)
B, H, L, D = 3, 2, 6, 8
S_q = 2

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H, S_q, D).astype(np.float32)
K = np.random.rand(B, H, L, D).astype(np.float32)
V = np.random.rand(B, H, L, D).astype(np.float32)
nonpad_kv_seqlen = np.array([4, 5, 6], dtype=np.int64)  # offsets [2, 3, 4]

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
    is_causal=1,
)

assert np.all(np.isfinite(Y)), "per-batch prefill output must be finite"

expect(
    node,
    inputs=[Q, K, V, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_causal_nonpad_batch_prefill",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_causal_nonpad_continued_prefill

"""Continued / chunked prefill (S_q=2) into a partially-filled static cache.

With ``nonpad_kv_seqlen = [4]`` and ``S_q = 2`` the bottom-right offset is
``4 - 2 = 2``: query 0 attends keys ``{0,1,2}`` and query 1 attends
``{0,1,2,3}``.  The old top-left alignment would mask everything past the
diagonal (``{0}`` and ``{0,1}``), so this test fails pre-fix.
"""
np.random.seed(1)
B, H, L, D = 1, 2, 4, 8
S_q = 2

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H, S_q, D).astype(np.float32)
K = np.random.rand(B, H, L, D).astype(np.float32)
V = np.random.rand(B, H, L, D).astype(np.float32)
nonpad_kv_seqlen = np.array([4], dtype=np.int64)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_causal_nonpad_continued_prefill",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_causal_nonpad_negative_offset_structural_empty

"""Negative bottom-right offset: structurally-empty early query rows -> zero.

This is the onnx-node twin of the ORT gtest
``Attention_Causal_NonPadKVSeqLen_StructuralEmptyRow_Zero`` /
``StructuralEmptyRows_Zero_CUDA``.  With ``nonpad_kv_seqlen = [2]`` and
``S_q = 4`` the bottom-right offset is ``2 - 4 = -2``: query row ``sq``
attends keys ``0..(sq - 2)``, so rows 0 and 1 have an empty key set.  Their
``softmax`` over an all-``-inf`` bias row is ``NaN``; the fully-masked-row
guard zeros those rows before the ``P @ V`` contraction so the output rows are
exactly ``0``, while rows 2 and 3 (attending keys ``{0}`` and ``{0,1}``) stay finite
and nonzero.  A ``nonpad_kv_seqlen[b] < q_sequence_length`` input is out of
the contract's intended use, but its result is still well-defined (zeroed
rows) rather than ``NaN``; this test pins that defined behavior.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(7)
B, H, L, D = 1, 2, 4, 8
S_q = 4

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H, S_q, D).astype(np.float32)
K = np.random.rand(B, H, L, D).astype(np.float32)
V = np.random.rand(B, H, L, D).astype(np.float32)
# offset = nonpad - S_q = 2 - 4 = -2 -> rows 0,1 structurally empty.
nonpad_kv_seqlen = np.array([2], dtype=np.int64)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
    is_causal=1,
)

# Structurally-empty early rows are exactly zero (not NaN); later rows finite.
assert np.all(np.isfinite(Y)), "all output rows must be finite"
assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), (
    "structurally-empty row 0 must be zero"
)
assert np.array_equal(Y[:, :, 1, :], np.zeros_like(Y[:, :, 1, :])), (
    "structurally-empty row 1 must be zero"
)
assert np.any(Y[:, :, 2, :] != 0) and np.any(Y[:, :, 3, :] != 0), (
    "rows with a non-empty key set must be nonzero"
)

expect(
    node,
    inputs=[Q, K, V, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_causal_nonpad_negative_offset_structural_empty",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_causal_with_past_and_present

"""Regression guard: internal (past_key) cache + is_causal.

This exercises the unchanged scalar bottom-right path (offset =
past_sequence_length).  Its golden output must remain identical to the
pre-fix behavior, proving the external-cache change does not touch the
past_key path.
"""
np.random.seed(2)
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
    is_causal=1,
)

past_sequence_length = 3
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 4, 8).astype(np.float32)
V = np.random.rand(2, 3, 4, 8).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    past_key=past_key,
    past_value=past_value,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_causal_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_diff_heads_mask4d_padded_kv

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 4).astype(np.float32)
nonpad_kv_seqlen = np.array([3, 4], dtype=np.int64)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_diff_heads_mask4d_padded_kv",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_gqa_causal_nonpad_decode

"""External/static-cache decode (S_q=1) with per-batch valid lengths.

K/V are the full static cache buffer; ``nonpad_kv_seqlen`` marks how many
leading keys are valid per batch.  With bottom-right (offset-aware) causal
masking the single decode query attends keys ``0..nonpad[b]-1``.  Under the
old top-left alignment it would attend only key 0, so this test fails
pre-fix and passes post-fix.
"""
np.random.seed(0)
B, H_q, H_kv, L, D = 2, 4, 2, 8, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H_q, 1, D).astype(np.float32)
K = np.random.rand(B, H_kv, L, D).astype(np.float32)
V = np.random.rand(B, H_kv, L, D).astype(np.float32)
# Batch 0 has all 8 keys valid, batch 1 only the first 5.
nonpad_kv_seqlen = np.array([8, 5], dtype=np.int64)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_gqa_causal_nonpad_decode",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_4d_gqa_causal_nonpad_decode_fp16

"""fp16 variant of the external-cache decode case (locks -inf dtype handling)."""
np.random.seed(0)
B, H_q, H_kv, L, D = 2, 4, 2, 8, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H_q, 1, D).astype(np.float16)
K = np.random.rand(B, H_kv, L, D).astype(np.float16)
V = np.random.rand(B, H_kv, L, D).astype(np.float16)
nonpad_kv_seqlen = np.array([8, 5], dtype=np.int64)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    nonpad_kv_seqlen=nonpad_kv_seqlen,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, nonpad_kv_seqlen],
    outputs=[Y],
    name="test_attention_4d_gqa_causal_nonpad_decode_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_attn_3d_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 1, 4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask_3d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_attn_3d_mask_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 1, 4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask_3d_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_attn_4d_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask_4d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_attn_4d_mask_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask_4d_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_attn_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_attn_mask_bool

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(bool)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask_bool",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_attn_mask_bool_4d

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 6).astype(bool)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_attn_mask_bool_4d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, is_causal=1)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_causal_boolmask_nan_robustness

"""Composed ``is_causal`` + boolean ``attn_mask`` NaN-robustness.

The causal frontier (lower-triangular here, offset 0) and the boolean
``attn_mask`` are intersected: a key is attended only if allowed by both.
This exercises two pre-fix NaN sources on the same forward pass:

* **Bug-1 (allowed cells stay finite).**  Query 0 is allowed key 0 by both
  the causal frontier (``{0}``) and the mask (``True`` at key 0).  The old
  ``(1 - attn_mask) * -inf`` conversion computes ``0 * -inf = NaN`` at that
  allowed cell, poisoning the row.  The select conversion
  ``where(attn_mask, 0, -inf)`` keeps it finite.
* **Bug-2 (fully-masked row -> 0).**  Query 1 is allowed keys ``{0, 1}`` by
  the causal frontier but the mask is ``False`` at both, so the combined
  constraint allows no key.  ``softmax`` of an all-``-inf`` row is ``NaN``;
  the fully-masked-row guard zeros it before the ``P @ V`` contraction so
  the output row is ``0``.

4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing
them would make the function body treat the input as 3D).
"""
np.random.seed(3)
B, H, S, D = 1, 2, 2, 8

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(B, H, S, D).astype(np.float32)
K = np.random.rand(B, H, S, D).astype(np.float32)
V = np.random.rand(B, H, S, D).astype(np.float32)
# Row 0: key 0 allowed (Bug-1 allowed cell). Row 1: no key allowed -> fully
# masked once intersected with the causal frontier (Bug-2 empty row).
attn_mask = np.array([[True, False], [False, False]], dtype=np.bool_)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    is_causal=1,
)

# Bug-1: allowed cells are finite (no NaN anywhere). Bug-2: the fully-masked
# query row is exactly zero, not NaN.
assert np.all(np.isfinite(Y)), "allowed cells must be finite (Bug-1)"
assert np.array_equal(Y[:, :, 1, :], np.zeros_like(Y[:, :, 1, :])), (
    "fully-masked row must be zero (Bug-2)"
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_causal_boolmask_nan_robustness",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

attention_diff_head_sizes

node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"])

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_diff_heads_sizes",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_attn_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_diff_heads_sizes_attn_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_diff_heads_sizes_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_scaled

scale = 1e-2
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, scale=scale)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_diff_heads_sizes_scaled",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_softcap

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    softcap=2.0,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    softcap=2.0,
)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_diff_heads_sizes_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_with_past_and_present

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_diff_heads_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_with_past_and_present_mask3D

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)
attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_diff_heads_with_past_and_present_mask3d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_diff_head_sizes_with_past_and_present_mask4D

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 10).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_diff_heads_with_past_and_present_mask4d",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_fp16

node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"])

Q = np.random.rand(2, 3, 4, 8).astype(np.float16)
K = np.random.rand(2, 3, 6, 8).astype(np.float16)
V = np.random.rand(2, 3, 6, 8).astype(np.float16)

Y, _, _, _ = _compute_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa

node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"])

Q = np.random.rand(2, 9, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_gqa",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa_attn_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
)

Q = np.random.rand(2, 9, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_gqa_attn_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    is_causal=1,
)

Q = np.random.rand(2, 9, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, is_causal=1)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_gqa_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa_scaled

scale = 1e-2
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
)

Q = np.random.rand(2, 9, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, scale=scale)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_gqa_scaled",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa_softcap

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    softcap=2.0,
)

Q = np.random.rand(2, 9, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, softcap=2.0)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_gqa_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa_with_past_and_present

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
)

past_sequence_length = 12
Q = np.random.rand(2, 9, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_gqa_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_gqa_with_past_and_present_fp16

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
)

past_sequence_length = 12
Q = np.random.rand(2, 9, 4, 8).astype(np.float16)
K = np.random.rand(2, 3, 6, 8).astype(np.float16)
V = np.random.rand(2, 3, 6, 8).astype(np.float16)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float16)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float16)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float16)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_gqa_with_past_and_present_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_scaled

scale = 1e-2
node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, scale=scale)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_scaled",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_softcap

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    softcap=2.0,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, _ = _compute_attention(Q, K, V, softcap=2.0)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_attention_4d_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_softcap_with_neginf_mask

"""Softcap + -inf mask: verifies softcap is applied BEFORE mask/bias.

If ordering were wrong (mask then softcap), tanh(-inf/softcap) = -1,
so softcap * tanh(-inf/softcap) = -softcap (finite).  That leaks
probability to masked positions.  With correct ordering (softcap then
mask), the -inf mask values survive to softmax and yield zero weight.
"""
np.random.seed(42)
B, H, S_q, S_kv, D = 1, 1, 4, 6, 8

Q = np.random.rand(B, H, S_q, D).astype(np.float32)
K = np.random.rand(B, H, S_kv, D).astype(np.float32)
V = np.random.rand(B, H, S_kv, D).astype(np.float32)

# All Q positions are blocked from KV positions 4 and 5.
attn_mask = np.zeros((S_q, S_kv), dtype=np.float32)
attn_mask[:, 4:] = -np.inf

softcap = 0.5

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    softcap=softcap,
)

Y, _, _, _ = _compute_attention(Q, K, V, attn_mask=attn_mask, softcap=softcap)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_softcap_neginf_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_softcap_with_neginf_mask_poison

"""Softcap + -inf mask + poison values at masked KV positions.

V has value 1000 at the masked positions (4 and 5).  With correct
ordering the output stays in [0, 1] because the mask zeros out those
positions.  With wrong ordering the output explodes (> 50), making
the failure obvious even with loose tolerances.
"""
np.random.seed(42)
B, H, S_q, S_kv, D = 1, 1, 4, 6, 8

Q = np.random.rand(B, H, S_q, D).astype(np.float32)
K = np.random.rand(B, H, S_kv, D).astype(np.float32)
V = np.random.rand(B, H, S_kv, D).astype(np.float32)

# Block all Q positions from KV positions 4 and 5.
attn_mask = np.zeros((S_q, S_kv), dtype=np.float32)
attn_mask[:, 4:] = -np.inf

# Poison: if attention leaks to masked positions, output >> 1.
V[:, :, 4:, :] = 1000.0

softcap = 0.5

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y"],
    softcap=softcap,
)

Y, _, _, _ = _compute_attention(Q, K, V, attn_mask=attn_mask, softcap=softcap)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y],
    name="test_attention_4d_softcap_neginf_mask_poison",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value"],
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, _ = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value],
    name="test_attention_4d_with_past_and_present",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present_qk_matmul

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_4d_with_past_and_present_qk_matmul",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present_qk_matmul_bias

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    qk_matmul_output_mode=2,
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    qk_matmul_output_mode=2,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_4d_with_past_and_present_qk_matmul_bias",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present_qk_matmul_bias_3d_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    qk_matmul_output_mode=2,
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    qk_matmul_output_mode=2,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present_qk_matmul_bias_3d_mask_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    qk_matmul_output_mode=2,
    is_causal=1,
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    qk_matmul_output_mode=2,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present_qk_matmul_bias_4d_mask

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    qk_matmul_output_mode=2,
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    qk_matmul_output_mode=2,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_past_and_present_qk_matmul_bias_4d_mask_causal

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
    outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
    qk_matmul_output_mode=2,
    is_causal=1,
)

past_sequence_length = 12
Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32)
past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)

Y, present_key, present_value, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    past_key=past_key,
    past_value=past_value,
    qk_matmul_output_mode=2,
    is_causal=1,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask, past_key, past_value],
    outputs=[Y, present_key, present_value, qk_matmul_output],
    name="test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_qk_matmul

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V"],
    outputs=["Y", "", "", "qk_matmul_output"],
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)

Y, _, _, qk_matmul_output = _compute_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y, qk_matmul_output],
    name="test_attention_4d_with_qk_matmul",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_qk_matmul_bias

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y", "", "", "qk_matmul_output"],
    qk_matmul_output_mode=2,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    qk_matmul_output_mode=2,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y, qk_matmul_output],
    name="test_attention_4d_with_qk_matmul_bias",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_qk_matmul_softcap

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y", "", "", "qk_matmul_output"],
    softcap=2.0,
    qk_matmul_output_mode=1,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    softcap=2.0,
    qk_matmul_output_mode=1,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y, qk_matmul_output],
    name="test_attention_4d_with_qk_matmul_softcap",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

attention_with_qk_matmul_softmax

node = onnx.helper.make_node(
    "Attention",
    inputs=["Q", "K", "V", "attn_mask"],
    outputs=["Y", "", "", "qk_matmul_output"],
    qk_matmul_output_mode=3,
)

Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
K = np.random.rand(2, 3, 6, 8).astype(np.float32)
V = np.random.rand(2, 3, 6, 8).astype(np.float32)
attn_mask = np.random.rand(4, 6).astype(np.float32)

Y, _, _, qk_matmul_output = _compute_attention(
    Q,
    K,
    V,
    attn_mask=attn_mask,
    qk_matmul_output_mode=3,
)

expect(
    node,
    inputs=[Q, K, V, attn_mask],
    outputs=[Y, qk_matmul_output],
    name="test_attention_4d_with_qk_matmul_softmax",
    opset_imports=[onnx.helper.make_opsetid("", 23)],
)

AveragePool

There are 17 test cases, listed as following:

averagepool_1d_default

"""input_shape: [1, 3, 32]
output_shape: [1, 3, 31]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2],
)
x = np.random.randn(1, 3, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = [2]
strides = [1]
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "AVG")

expect(node, inputs=[x], outputs=[y], name="test_averagepool_1d_default")

averagepool_2d_ceil

"""input_shape: [1, 1, 4, 4]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    strides=[2, 2],
    ceil_mode=True,
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[6, 7.5], [12, 13.5]]]]).astype(np.float32)

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_ceil")

averagepool_2d_ceil_last_window_starts_on_pad

"""input_shape: [1, 3, 2, 2]
output_shape: [1, 3, 1, 1]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    strides=[3, 3],
    pads=[1, 1, 1, 1],
    ceil_mode=True,
    count_include_pad=1,
)
x = np.array(
    [
        [
            [[0.8580, 0.0786], [0.2692, 0.1537]],
            [[0.8816, 0.4353], [0.5772, 0.6623]],
            [[0.9067, 0.9483], [0.5970, 0.7630]],
        ]
    ]
).astype(np.float32)
y = np.array([[[[0.1511]], [[0.2841]], [[0.3572]]]]).astype(np.float32)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_averagepool_2d_ceil_last_window_starts_on_pad",
)

averagepool_2d_default

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 31, 31]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = (2, 2)
strides = (1, 1)
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "AVG")

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_default")

averagepool_2d_dilations

"""input_shape: [1, 1, 4, 4]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    strides=[1, 1],
    dilations=[2, 2],
    ceil_mode=True,
)

# input shape: [1, 1, 4, 4]
x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)

y = np.array([[[[6, 7], [10, 11]]]]).astype(np.float32)

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_dilations")

averagepool_2d_pads

"""input_shape: [1, 3, 28, 28]
output_shape: [1, 3, 30, 30]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[2, 2, 2, 2],
)
x = np.random.randn(1, 3, 28, 28).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (3, 3)
strides = (1, 1)
pad_bottom = 2
pad_top = 2
pad_right = 2
pad_left = 2
pads = [pad_top, pad_left, pad_bottom, pad_right]
out_shape, extra_pads = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides, ceil_mode=False
)
padded = np.pad(
    x,
    (
        (0, 0),
        (0, 0),
        (extra_pads[0], extra_pads[2]),
        (extra_pads[1], extra_pads[3]),
    ),
    mode="constant",
    constant_values=np.nan,
)
y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "AVG",
    pads_required=extra_pads,
    pads=pads,
)

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_pads")

averagepool_2d_pads_count_include_pad

"""input_shape: [1, 3, 28, 28]
output_shape: [1, 3, 30, 30]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[2, 2, 2, 2],
    count_include_pad=1,
)
x = np.random.randn(1, 3, 28, 28).astype(np.float32)
x_shape = np.shape(x)
dilations = (1, 1)
kernel_shape = (3, 3)
strides = (1, 1)
pad_bottom = 2
pad_top = 2
pad_right = 2
pad_left = 2
pads = [pad_top, pad_left, pad_bottom, pad_right]
out_shape, extra_pads = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides, dilations, ceil_mode=False
)
padded = np.pad(
    x,
    (
        (0, 0),
        (0, 0),
        (extra_pads[0], extra_pads[2]),
        (extra_pads[1], extra_pads[3]),
    ),
    mode="constant",
    constant_values=0,
)
y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "AVG",
    pads_required=extra_pads,
    pads=pads,
    count_include_pad=1,
)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_averagepool_2d_pads_count_include_pad",
)

averagepool_2d_precomputed_pads

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 5, 5]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[5, 5],
    pads=[2, 2, 2, 2],
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array(
    [
        [
            [
                [7, 7.5, 8, 8.5, 9],
                [9.5, 10, 10.5, 11, 11.5],
                [12, 12.5, 13, 13.5, 14],
                [14.5, 15, 15.5, 16, 16.5],
                [17, 17.5, 18, 18.5, 19],
            ]
        ]
    ]
).astype(np.float32)

expect(
    node, inputs=[x], outputs=[y], name="test_averagepool_2d_precomputed_pads"
)

averagepool_2d_precomputed_pads_count_include_pad

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 5, 5]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[5, 5],
    pads=[2, 2, 2, 2],
    count_include_pad=1,
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array(
    [
        [
            [
                [2.5200, 3.6000, 4.8000, 4.0800, 3.2400],
                [4.5600, 6.4000, 8.4000, 7.0400, 5.5200],
                [7.2000, 10.0000, 13.0000, 10.8000, 8.4000],
                [6.9600, 9.6000, 12.4000, 10.2400, 7.9200],
                [6.1200, 8.4000, 10.8000, 8.8800, 6.8400],
            ]
        ]
    ]
).astype(np.float32)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_averagepool_2d_precomputed_pads_count_include_pad",
)

averagepool_2d_precomputed_same_upper

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 3, 3]
pad_shape: [2, 2] -> [1, 1, 1, 1] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    strides=[2, 2],
    auto_pad="SAME_UPPER",
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[4, 5.5, 7], [11.5, 13, 14.5], [19, 20.5, 22]]]]).astype(
    np.float32
)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_averagepool_2d_precomputed_same_upper",
)

averagepool_2d_precomputed_strides

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    strides=[2, 2],
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[4, 6], [14, 16]]]]).astype(np.float32)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_averagepool_2d_precomputed_strides",
)

averagepool_2d_same_lower

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 32, 32]
pad_shape: [1, 1] -> [1, 0, 1, 0] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    auto_pad="SAME_LOWER",
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (2, 2)
strides = (1, 1)
out_shape = get_output_shape_auto_pad(
    "SAME_LOWER", x_shape[2:], kernel_shape, strides
)
pad_shape = get_pad_shape(
    "SAME_LOWER", x_shape[2:], kernel_shape, strides, out_shape
)
pad_bottom = pad_shape[0] // 2
pad_top = pad_shape[0] - pad_bottom
pad_right = pad_shape[1] // 2
pad_left = pad_shape[1] - pad_right
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=np.nan,
)
pads = (pad_top, pad_left, pad_bottom, pad_right)
y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "AVG",
    pads_required=pads,
    pads=pads,
)

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_same_lower")

averagepool_2d_same_upper

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 32, 32]
pad_shape: [1, 1] -> [0, 1, 0, 1] by axis
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    auto_pad="SAME_UPPER",
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (2, 2)
strides = (1, 1)
out_shape = get_output_shape_auto_pad(
    "SAME_UPPER", x_shape[2:], kernel_shape, strides
)
pad_shape = get_pad_shape(
    "SAME_UPPER", x_shape[2:], kernel_shape, strides, out_shape
)
pad_top = pad_shape[0] // 2
pad_bottom = pad_shape[0] - pad_top
pad_left = pad_shape[1] // 2
pad_right = pad_shape[1] - pad_left
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=np.nan,
)
pads = (pad_top, pad_left, pad_bottom, pad_right)
y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "AVG",
    pads_required=pads,
    pads=pads,
)

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_same_upper")

averagepool_2d_strides

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 10, 10]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[5, 5],
    strides=[3, 3],
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (5, 5)
strides = (3, 3)
out_shape, pads = get_output_shape_explicit_padding(
    None, x_shape[2:], kernel_shape, strides, ceil_mode=False
)
padded = x
y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "AVG",
    pads_required=pads,
    pads=None,
)

expect(node, inputs=[x], outputs=[y], name="test_averagepool_2d_strides")

averagepool_3d_default

"""input_shape: [1, 3, 32, 32, 32]
output_shape: [1, 3, 31, 31, 31]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2, 2],
)
x = np.random.randn(1, 3, 32, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = [2, 2, 2]
strides = [1, 1, 1]
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "AVG")

expect(node, inputs=[x], outputs=[y], name="test_averagepool_3d_default")

averagepool_3d_dilations

"""input_shape: [1, 1, 4, 4]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "AveragePool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2, 2],
    strides=[1, 1, 1],
    dilations=[2, 2, 2],
    ceil_mode=True,
)

# input shape: [1, 1, 4, 4, 4]
x = np.array(
    [
        [
            [
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
            ]
        ]
    ]
).astype(np.float32)

y = np.array([[[[[6, 7], [10, 11]], [[6, 7], [10, 11]]]]]).astype(np.float32)

expect(
    node, inputs=[x], outputs=[y], name="test_averagepool_3d_dilations_small"
)

averagepool_3d_dilations_large

x_shape = (32, 32, 32)
dilations = (2, 2, 2)
kernel_shape = (5, 5, 5)
strides = (3, 3, 3)
count_include_pad = 0

for count_include_pad in (0, 1):
    for ceil_mode in (True, False):
        node = onnx.helper.make_node(
            "AveragePool",
            inputs=["x"],
            outputs=["y"],
            kernel_shape=kernel_shape,
            strides=strides,
            dilations=dilations,
            count_include_pad=count_include_pad,
            ceil_mode=ceil_mode,
        )

        x = np.random.randn(1, 1, *x_shape).astype(np.float32)
        out_shape, extra_pads = get_output_shape_explicit_padding(
            None,
            x_shape,
            kernel_shape,
            strides,
            dilations=dilations,
            ceil_mode=ceil_mode,
        )
        padded = np.pad(
            x,
            (
                (0, 0),
                (0, 0),
                (extra_pads[0], extra_pads[3]),
                (extra_pads[1], extra_pads[4]),
                (extra_pads[2], extra_pads[5]),
            ),
            mode="constant",
            constant_values=0 if count_include_pad == 1 else np.nan,
        )
        y = pool(
            padded,
            (1, 1, *x_shape),
            kernel_shape,
            strides,
            out_shape,
            "AVG",
            pads_required=extra_pads,
            pads=None,
            dilations=dilations,
            count_include_pad=count_include_pad,
        )

        test_name = f"test_averagepool_3d_dilations_large_count_include_pad_is_{count_include_pad}_ceil_mode_is_{ceil_mode}"
        expect(node, inputs=[x], outputs=[y], name=test_name)

BatchNormalization

There are 2 test cases, listed as following:

batchnormalization

# input size: (2, 3, 4, 5)
x = np.random.randn(2, 3, 4, 5).astype(np.float32)
s = np.random.randn(3).astype(np.float32)
bias = np.random.randn(3).astype(np.float32)
mean = np.random.randn(3).astype(np.float32)
var = np.random.rand(3).astype(np.float32)
y = _batchnorm_test_mode(x, s, bias, mean, var).astype(np.float32)

node = onnx.helper.make_node(
    "BatchNormalization",
    inputs=["x", "s", "bias", "mean", "var"],
    outputs=["y"],
)

# output size: (2, 3, 4, 5)
expect(
    node,
    inputs=[x, s, bias, mean, var],
    outputs=[y],
    name="test_batchnorm_example",
)

# input size: (2, 3, 4, 5)
x = np.random.randn(2, 3, 4, 5).astype(np.float32)
s = np.random.randn(3).astype(np.float32)
bias = np.random.randn(3).astype(np.float32)
mean = np.random.randn(3).astype(np.float32)
var = np.random.rand(3).astype(np.float32)
epsilon = 1e-2
y = _batchnorm_test_mode(x, s, bias, mean, var, epsilon).astype(np.float32)

node = onnx.helper.make_node(
    "BatchNormalization",
    inputs=["x", "s", "bias", "mean", "var"],
    outputs=["y"],
    epsilon=epsilon,
)

# output size: (2, 3, 4, 5)
expect(
    node,
    inputs=[x, s, bias, mean, var],
    outputs=[y],
    name="test_batchnorm_epsilon",
)

train

# input size: (2, 3, 4, 5)
x = np.random.randn(2, 3, 4, 5).astype(np.float32)
s = np.random.randn(3).astype(np.float32)
bias = np.random.randn(3).astype(np.float32)
mean = np.random.randn(3).astype(np.float32)
var = np.random.rand(3).astype(np.float32)
# using np.bool(1) while generating test data with "'bool' object has no attribute 'dtype'"
# working around by using np.byte(1).astype(bool)
training_mode = 1
y, output_mean, output_var = _batchnorm_training_mode(x, s, bias, mean, var)

node = onnx.helper.make_node(
    "BatchNormalization",
    inputs=["x", "s", "bias", "mean", "var"],
    outputs=["y", "output_mean", "output_var"],
    training_mode=training_mode,
)

# output size: (2, 3, 4, 5)
expect(
    node,
    inputs=[x, s, bias, mean, var],
    outputs=[y, output_mean, output_var],
    name="test_batchnorm_example_training_mode",
)

# input size: (2, 3, 4, 5)
x = np.random.randn(2, 3, 4, 5).astype(np.float32)
s = np.random.randn(3).astype(np.float32)
bias = np.random.randn(3).astype(np.float32)
mean = np.random.randn(3).astype(np.float32)
var = np.random.rand(3).astype(np.float32)
training_mode = 1
momentum = 0.9
epsilon = 1e-2
y, output_mean, output_var = _batchnorm_training_mode(
    x, s, bias, mean, var, momentum, epsilon
)

node = onnx.helper.make_node(
    "BatchNormalization",
    inputs=["x", "s", "bias", "mean", "var"],
    outputs=["y", "output_mean", "output_var"],
    epsilon=epsilon,
    training_mode=training_mode,
)

# output size: (2, 3, 4, 5)
expect(
    node,
    inputs=[x, s, bias, mean, var],
    outputs=[y, output_mean, output_var],
    name="test_batchnorm_epsilon_training_mode",
)

Bernoulli

There are 3 test cases, listed as following:

bernoulli_with_dtype

node = onnx.helper.make_node(
    "Bernoulli",
    inputs=["x"],
    outputs=["y"],
    dtype=onnx.TensorProto.DOUBLE,
)

x = np.random.uniform(0.0, 1.0, 10).astype(np.float32)
y = bernoulli_reference_implementation(x, float)
expect(node, inputs=[x], outputs=[y], name="test_bernoulli_double")

bernoulli_with_seed

seed = float(0)
node = onnx.helper.make_node(
    "Bernoulli",
    inputs=["x"],
    outputs=["y"],
    seed=seed,
)

x = np.random.uniform(0.0, 1.0, 10).astype(np.float32)
y = bernoulli_reference_implementation(x, np.float32)
expect(node, inputs=[x], outputs=[y], name="test_bernoulli_seed")

bernoulli_without_dtype

node = onnx.helper.make_node(
    "Bernoulli",
    inputs=["x"],
    outputs=["y"],
)

x = np.random.uniform(0.0, 1.0, 10).astype(float)
y = bernoulli_reference_implementation(x, float)
expect(node, inputs=[x], outputs=[y], name="test_bernoulli")

BitCast

There are 10 test cases, listed as following:

bitcast_2d_float32_to_int32

"""Test bitcasting 2D array from float32 to int32."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.INT32,
)
x = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
y = x.view(np.int32)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_2d_float32_to_int32")

bitcast_bool_to_uint8

"""Test bitcasting from bool to uint8 (same size)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.UINT8,
)
x = np.array([True, False, True, False], dtype=np.bool_)
y = x.view(np.uint8)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_bool_to_uint8")

bitcast_float32_to_int32

"""Test bitcasting from float32 to int32 (same size)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.INT32,
)
x = np.array([1.0, -2.5, 3.75], dtype=np.float32)
y = x.view(np.int32)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_float32_to_int32")

bitcast_float64_to_int64

"""Test bitcasting from float64 to int64 (same size)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.INT64,
)
x = np.array([1.0, -2.5, 3.75], dtype=np.float64)
y = x.view(np.int64)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_float64_to_int64")

bitcast_int32_to_float32

"""Test bitcasting from int32 to float32 (same size)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.FLOAT,
)
x = np.array([1065353216, -1071644672, 1081081856], dtype=np.int32)
y = x.view(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_int32_to_float32")

bitcast_int64_to_float64

"""Test bitcasting from int64 to float64 (same size)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.DOUBLE,
)
x = np.array(
    [4607182418800017408, -4611686018427387904, 4614256656552045184],
    dtype=np.int64,
)
y = x.view(np.float64)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_int64_to_float64")

bitcast_int8_to_uint8

"""Test bitcasting from int8 to uint8 (same size, different signedness)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.UINT8,
)
x = np.array([-1, -128, 127, 0], dtype=np.int8)
y = x.view(np.uint8)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_int8_to_uint8")

bitcast_scalar_float32_to_int32

"""Test bitcasting scalar from float32 to int32."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.INT32,
)
x = np.array(1.0, dtype=np.float32)
y = x.view(np.int32)
expect(
    node, inputs=[x], outputs=[y], name="test_bitcast_scalar_float32_to_int32"
)

bitcast_uint16_to_int16

"""Test bitcasting from uint16 to int16 (same size, different signedness)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.INT16,
)
x = np.array([1, 32768, 65535], dtype=np.uint16)
y = x.view(np.int16)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_uint16_to_int16")

bitcast_uint32_to_int32

"""Test bitcasting from uint32 to int32 (same size, different signedness)."""
node = onnx.helper.make_node(
    "BitCast",
    inputs=["x"],
    outputs=["y"],
    to=onnx.TensorProto.INT32,
)
x = np.array([4294967295, 2147483648, 2147483647], dtype=np.uint32)
y = x.view(np.int32)
expect(node, inputs=[x], outputs=[y], name="test_bitcast_uint32_to_int32")

BitShift

There are 8 test cases, listed as following:

left_unit16

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="LEFT"
)

x = np.array([16, 4, 1]).astype(np.uint16)
y = np.array([1, 2, 3]).astype(np.uint16)
z = x << y  # expected output [32, 16, 8]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_left_uint16")

left_unit32

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="LEFT"
)

x = np.array([16, 4, 1]).astype(np.uint32)
y = np.array([1, 2, 3]).astype(np.uint32)
z = x << y  # expected output [32, 16, 8]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_left_uint32")

left_unit64

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="LEFT"
)

x = np.array([16, 4, 1]).astype(np.uint64)
y = np.array([1, 2, 3]).astype(np.uint64)
z = x << y  # expected output [32, 16, 8]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_left_uint64")

left_unit8

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="LEFT"
)

x = np.array([16, 4, 1]).astype(np.uint8)
y = np.array([1, 2, 3]).astype(np.uint8)
z = x << y  # expected output [32, 16, 8]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_left_uint8")

right_unit16

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="RIGHT"
)

x = np.array([16, 4, 1]).astype(np.uint16)
y = np.array([1, 2, 3]).astype(np.uint16)
z = x >> y  # expected output [8, 1, 0]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_right_uint16")

right_unit32

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="RIGHT"
)

x = np.array([16, 4, 1]).astype(np.uint32)
y = np.array([1, 2, 3]).astype(np.uint32)
z = x >> y  # expected output [8, 1, 0]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_right_uint32")

right_unit64

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="RIGHT"
)

x = np.array([16, 4, 1]).astype(np.uint64)
y = np.array([1, 2, 3]).astype(np.uint64)
z = x >> y  # expected output [8, 1, 0]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_right_uint64")

right_unit8

node = onnx.helper.make_node(
    "BitShift", inputs=["x", "y"], outputs=["z"], direction="RIGHT"
)

x = np.array([16, 4, 1]).astype(np.uint8)
y = np.array([1, 2, 3]).astype(np.uint8)
z = x >> y  # expected output [8, 1, 0]
expect(node, inputs=[x, y], outputs=[z], name="test_bitshift_right_uint8")

BitwiseAnd

There are 2 test cases, listed as following:

bitwiseand

node = onnx.helper.make_node(
    "BitwiseAnd",
    inputs=["x", "y"],
    outputs=["bitwiseand"],
)

# 2d
x = create_random_int((3, 4), np.int32)
y = create_random_int((3, 4), np.int32)
z = np.bitwise_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_and_i32_2d")

# 3d
x = create_random_int((3, 4, 5), np.int16)
y = create_random_int((3, 4, 5), np.int16)
z = np.bitwise_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_and_i16_3d")

bitwiseand_broadcast

node = onnx.helper.make_node(
    "BitwiseAnd",
    inputs=["x", "y"],
    outputs=["bitwiseand"],
)

# 3d vs 1d
x = create_random_int((3, 4, 5), np.uint64)
y = create_random_int((5,), np.uint64)
z = np.bitwise_and(x, y)
expect(
    node, inputs=[x, y], outputs=[z], name="test_bitwise_and_ui64_bcast_3v1d"
)

# 4d vs 3d
x = create_random_int((3, 4, 5, 6), np.uint8)
y = create_random_int((4, 5, 6), np.uint8)
z = np.bitwise_and(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_and_ui8_bcast_4v3d")

BitwiseNot

There are 1 test cases, listed as following:

bitwisenot

node = onnx.helper.make_node(
    "BitwiseNot",
    inputs=["x"],
    outputs=["bitwise_not"],
)

# 2d
x = create_random_int((3, 4), np.int32)
y = np.bitwise_not(x)
expect(node, inputs=[x], outputs=[y], name="test_bitwise_not_2d")

# 3d
x = create_random_int((3, 4, 5), np.uint16)
y = np.bitwise_not(x)
expect(node, inputs=[x], outputs=[y], name="test_bitwise_not_3d")

# 4d
x = create_random_int((3, 4, 5, 6), np.uint8)
y = np.bitwise_not(x)
expect(node, inputs=[x], outputs=[y], name="test_bitwise_not_4d")

BitwiseOr

There are 2 test cases, listed as following:

bitwiseor

node = onnx.helper.make_node(
    "BitwiseOr",
    inputs=["x", "y"],
    outputs=["bitwiseor"],
)
# 2d
x = create_random_int((3, 4), np.int32)
y = create_random_int((3, 4), np.int32)
z = np.bitwise_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_or_i32_2d")

# 4d
x = create_random_int((3, 4, 5, 6), np.int8)
y = create_random_int((3, 4, 5, 6), np.int8)
z = np.bitwise_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_or_i16_4d")

bitwiseor_broadcast

node = onnx.helper.make_node(
    "BitwiseOr",
    inputs=["x", "y"],
    outputs=["bitwiseor"],
)

# 3d vs 1d
x = create_random_int((3, 4, 5), np.uint64)
y = create_random_int((5,), np.uint64)
z = np.bitwise_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_or_ui64_bcast_3v1d")

# 4d vs 3d
x = create_random_int((3, 4, 5, 6), np.uint8)
y = create_random_int((4, 5, 6), np.uint8)
z = np.bitwise_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_or_ui8_bcast_4v3d")

BitwiseXor

There are 2 test cases, listed as following:

bitwiseor_broadcast

node = onnx.helper.make_node(
    "BitwiseXor",
    inputs=["x", "y"],
    outputs=["bitwisexor"],
)

# 3d vs 1d
x = create_random_int((3, 4, 5), np.uint64)
y = create_random_int((5,), np.uint64)
z = np.bitwise_xor(x, y)
expect(
    node, inputs=[x, y], outputs=[z], name="test_bitwise_xor_ui64_bcast_3v1d"
)

# 4d vs 3d
x = create_random_int((3, 4, 5, 6), np.uint8)
y = create_random_int((4, 5, 6), np.uint8)
z = np.bitwise_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_xor_ui8_bcast_4v3d")

bitwisexor

node = onnx.helper.make_node(
    "BitwiseXor",
    inputs=["x", "y"],
    outputs=["bitwisexor"],
)

# 2d
x = create_random_int((3, 4), np.int32)
y = create_random_int((3, 4), np.int32)
z = np.bitwise_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_xor_i32_2d")

# 3d
x = create_random_int((3, 4, 5), np.int16)
y = create_random_int((3, 4, 5), np.int16)
z = np.bitwise_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_bitwise_xor_i16_3d")

BlackmanWindow

There are 1 test cases, listed as following:

blackmanwindow

# Test periodic window
node = onnx.helper.make_node(
    "BlackmanWindow",
    inputs=["x"],
    outputs=["y"],
)
size = np.int32(10)
a0 = 0.42
a1 = -0.5
a2 = 0.08
y = a0
y += a1 * np.cos(2 * np.pi * np.arange(0, size, 1, dtype=np.float32) / size)
y += a2 * np.cos(4 * np.pi * np.arange(0, size, 1, dtype=np.float32) / size)
expect(
    node,
    inputs=[size],
    outputs=[y.astype(np.float32)],
    name="test_blackmanwindow",
)

# Test symmetric window
node = onnx.helper.make_node(
    "BlackmanWindow", inputs=["x"], outputs=["y"], periodic=0
)
size = np.int32(10)
a0 = 0.42
a1 = -0.5
a2 = 0.08
y = a0
y += a1 * np.cos(
    2 * np.pi * np.arange(0, size, 1, dtype=np.float32) / (size - 1)
)
y += a2 * np.cos(
    4 * np.pi * np.arange(0, size, 1, dtype=np.float32) / (size - 1)
)
expect(
    node,
    inputs=[size],
    outputs=[y.astype(np.float32)],
    name="test_blackmanwindow_symmetric",
)

Cast

There are 3 test cases, listed as following:

cast

test_cases = [
    ("FLOAT", "FLOAT16"),
    ("FLOAT", "DOUBLE"),
    ("FLOAT16", "FLOAT"),
    ("FLOAT16", "DOUBLE"),
    ("DOUBLE", "FLOAT"),
    ("DOUBLE", "FLOAT16"),
    ("FLOAT", "BFLOAT16"),
    ("BFLOAT16", "FLOAT"),
    ("FLOAT", "FLOAT8E4M3FN"),
    ("FLOAT16", "FLOAT8E4M3FN"),
    ("FLOAT", "FLOAT8E4M3FNUZ"),
    ("FLOAT16", "FLOAT8E4M3FNUZ"),
    ("FLOAT8E4M3FN", "FLOAT"),
    ("FLOAT8E4M3FN", "FLOAT16"),
    ("FLOAT8E4M3FNUZ", "FLOAT"),
    ("FLOAT8E4M3FNUZ", "FLOAT16"),
    ("FLOAT", "FLOAT8E5M2"),
    ("FLOAT16", "FLOAT8E5M2"),
    ("FLOAT", "FLOAT8E5M2FNUZ"),
    ("FLOAT16", "FLOAT8E5M2FNUZ"),
    ("FLOAT8E5M2", "FLOAT"),
    ("FLOAT8E5M2", "FLOAT16"),
    ("FLOAT8E5M2FNUZ", "FLOAT"),
    ("FLOAT8E5M2FNUZ", "FLOAT16"),
    ("FLOAT", "UINT4"),
    ("FLOAT16", "UINT4"),
    ("FLOAT", "INT4"),
    ("FLOAT16", "INT4"),
    ("UINT4", "FLOAT"),
    ("UINT4", "FLOAT16"),
    ("UINT4", "UINT8"),
    ("INT4", "FLOAT"),
    ("INT4", "FLOAT16"),
    ("INT4", "INT8"),
    ("FLOAT4E2M1", "FLOAT"),
    ("FLOAT4E2M1", "FLOAT16"),
    ("FLOAT", "FLOAT4E2M1"),
    ("FLOAT16", "FLOAT4E2M1"),
    ("FLOAT", "UINT2"),
    ("FLOAT16", "UINT2"),
    ("FLOAT", "INT2"),
    ("FLOAT16", "INT2"),
    ("UINT2", "FLOAT"),
    ("UINT2", "FLOAT16"),
    ("UINT2", "UINT8"),
    ("INT2", "FLOAT"),
    ("INT2", "FLOAT16"),
    ("INT2", "INT8"),
]

for from_type, to_type in test_cases:
    if from_type == to_type:
        # Skip cases where from_type and to_type are the same
        continue
    from_dtype = getattr(TensorProto, from_type)
    to_dtype = getattr(TensorProto, to_type)
    from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
    to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)

    if from_type == "BFLOAT16" or to_type == "BFLOAT16":
        np_fp32 = np.array(
            [
                "0.47892547",
                "0.48033667",
                "0.49968487",
                "0.81910545",
                "0.47031248",
                "0.816468",
                "0.21087195",
                "0.7229038",
                "NaN",
                "INF",
                "+INF",
                "-INF",
            ],
            dtype=np.float32,
        )
        input_shape = (3, 4)

    elif from_type in F8_TYPES or to_type in F8_TYPES:
        np_fp32 = np.array(
            [
                "0.47892547",
                "0.48033667",
                "0.49968487",
                "0.81910545",
                "0.47031248",
                "0.7229038",
                "1000000",
                "1e-7",
                "NaN",
                "INF",
                "+INF",
                "-INF",
                "-0.0000001",
                "0.0000001",
                "-1000000",
            ],
            dtype=np.float32,
        )
        input_shape = (3, 5)
    elif from_type in ("UINT4", "INT4") or to_type in ("UINT4", "INT4"):
        np_fp32 = np.arange(-9, 16).astype(np.float32)
        input_shape = (5, 5)
    elif from_type in ("UINT2", "INT2") or to_type in ("UINT2", "INT2"):
        np_fp32 = np.arange(-3, 4).astype(np.float32)
        input_shape = (7, 1)
    elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
        np_fp32 = np.array(
            [
                "0.48",
                "0.25",
                "1.05",
                "-3.5",
                "-8",
                "9",
                "1000000",
                "1e-7",
                "NaN",
                "INF",
                "+INF",
                "-INF",
                "-4",
                "0.01",
                "-0.0",
            ],
            dtype=np.float32,
        )
        input_shape = (3, 5)

    else:
        np_fp32 = np.array(
            [
                "0.47892547",
                "0.48033667",
                "0.49968487",
                "0.81910545",
                "0.47031248",
                "0.816468",
                "0.21087195",
                "0.7229038",
                "NaN",
                "INF",
                "+INF",
                "-INF",
            ],
            dtype=np.float32,
        ).reshape([3, 4])
        input_shape = (3, 4)

    if from_type in F8_TYPES:
        np_from = onnx.numpy_helper.saturate_cast(np_fp32, from_np_dtype)
        input = make_tensor(
            "input",
            from_dtype,
            input_shape,
            vals=np_from,
            raw=True,
        )
    elif from_type in FOUR_BIT_TYPES:
        np_from = np_fp32.astype(from_np_dtype)
        packed = onnx.numpy_helper._pack_4bitx2(np_from)
        # No byteswap needed on big-endian machines as _pack_4bitx2()
        # returns a numpy array with uint8 datatype.
        input = make_tensor(
            "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    elif from_type in TWO_BIT_TYPES:
        np_from = np_fp32.astype(from_np_dtype)
        packed = onnx.numpy_helper._pack_2bitx4(np_from)
        input = make_tensor(
            "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    else:
        np_from = np_fp32.astype(from_np_dtype)
        input = make_tensor(
            "input", from_dtype, input_shape, vals=np_from, raw=True
        )

    if to_type in F8_TYPES:
        output = make_tensor(
            "output",
            to_dtype,
            input_shape,
            vals=onnx.numpy_helper.saturate_cast(np_from, to_np_dtype),
            raw=True,
        )
    elif to_type in FOUR_BIT_TYPES:
        packed = onnx.numpy_helper._pack_4bitx2(np_from.astype(to_np_dtype))
        # No byteswap needed on big-endian machines as _pack_4bitx2()
        # returns a numpy array with uint8 datatype.
        output = make_tensor(
            "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    elif to_type in TWO_BIT_TYPES:
        packed = onnx.numpy_helper._pack_2bitx4(np_from.astype(to_np_dtype))
        output = make_tensor(
            "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    else:
        output = make_tensor(
            "output",
            to_dtype,
            input_shape,
            vals=np_from.astype(to_np_dtype),
            raw=True,
        )

    node = onnx.helper.make_node(
        "Cast",
        inputs=["input"],
        outputs=["output"],
        to=to_dtype,
    )
    expect(
        node,
        inputs=[input],
        outputs=[output],
        name="test_cast_" + from_type + "_to_" + to_type,
    )

e8m0

np_fp32 = np.array(
    [
        "0.0",
        "0.124",
        "0.25",
        "0.5",
        "1.1",
        "2.0",
        "4.0",
        "8.0",
    ],
    dtype=np.float32,
)
test_cases = [
    ("FLOAT", "FLOAT8E8M0"),
    ("FLOAT16", "FLOAT8E8M0"),
    ("FLOAT8E8M0", "FLOAT"),
    ("FLOAT8E8M0", "FLOAT16"),
]
for from_type, to_type in test_cases:
    if from_type == "FLOAT":
        input_np = np_fp32
        output_np = to_float8e8m0(np_fp32)
    elif from_type == "FLOAT16":
        input_np = np_fp32.astype(np.float16)
        output_np = to_float8e8m0(input_np)
    elif from_type == "FLOAT8E8M0":
        input_np = to_float8e8m0(np_fp32)
        if to_type == "FLOAT":
            output_np = input_np.astype(np.float32)
        elif to_type == "FLOAT16":
            output_np = input_np.astype(np.float16)
        else:
            raise ValueError(
                f"Conversion from {from_type} to {to_type} is not tested."
            )
    else:
        raise ValueError(
            f"Conversion from {from_type} to {to_type} is not tested."
        )
    input = make_tensor(
        "input",
        getattr(TensorProto, from_type),
        [2, 4],
        input_np,
        raw=True,
    )
    output = make_tensor(
        "output",
        getattr(TensorProto, to_type),
        [2, 4],
        output_np,
        raw=True,
    )
    if to_type == "FLOAT8E8M0":
        node = onnx.helper.make_node(
            "Cast",
            inputs=["input"],
            outputs=["output"],
            to=getattr(TensorProto, to_type),
            saturate=1,
            round_mode="up",
        )
    else:
        node = onnx.helper.make_node(
            "Cast",
            inputs=["input"],
            outputs=["output"],
            to=getattr(TensorProto, to_type),
        )

    expect(
        node,
        inputs=[input],
        outputs=[output],
        name="test_cast_e8m0_" + from_type + "_to_" + to_type,
    )

saturate_false

test_cases = itertools.product(
    [
        "FLOAT",
        "FLOAT16",
    ],
    [
        "FLOAT8E4M3FN",
        "FLOAT8E4M3FNUZ",
        "FLOAT8E5M2",
        "FLOAT8E5M2FNUZ",
    ],
)
input_shape = (3, 5)
for from_type, to_type in test_cases:
    from_dtype = getattr(TensorProto, from_type)
    to_dtype = getattr(TensorProto, to_type)
    from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
    to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)
    np_fp32 = np.array(
        [
            "0.47892547",
            "0.48033667",
            "0.49968487",
            "0.81910545",
            "0.47031248",
            "0.7229038",
            "1000000",
            "1e-7",
            "NaN",
            "INF",
            "+INF",
            "-INF",
            "-0.0000001",
            "0.0000001",
            "-1000000",
        ],
        dtype=np.float32,
    )

    input = make_tensor(
        "input",
        from_dtype,
        input_shape,
        vals=np_fp32.astype(from_np_dtype),
        raw=True,
    )
    output = make_tensor(
        "output",
        to_dtype,
        input_shape,
        vals=np_fp32.astype(from_np_dtype).astype(to_np_dtype),
        raw=True,
    )

    node = onnx.helper.make_node(
        "Cast",
        inputs=["input"],
        outputs=["output"],
        to=to_dtype,
        saturate=0,
    )
    expect(
        node,
        inputs=[input],
        outputs=[output],
        name="test_cast_no_saturate_" + from_type + "_to_" + to_type,
    )

CastLike

There are 2 test cases, listed as following:

castlike

test_cases = [
    ("FLOAT", "FLOAT16"),
    ("FLOAT", "DOUBLE"),
    ("FLOAT16", "FLOAT"),
    ("FLOAT16", "DOUBLE"),
    ("DOUBLE", "FLOAT"),
    ("DOUBLE", "FLOAT16"),
    ("FLOAT", "BFLOAT16"),
    ("BFLOAT16", "FLOAT"),
    ("FLOAT", "FLOAT8E4M3FN"),
    ("FLOAT16", "FLOAT8E4M3FN"),
    ("FLOAT", "FLOAT8E4M3FNUZ"),
    ("FLOAT16", "FLOAT8E4M3FNUZ"),
    ("FLOAT8E4M3FN", "FLOAT"),
    ("FLOAT8E4M3FN", "FLOAT16"),
    ("FLOAT8E4M3FNUZ", "FLOAT"),
    ("FLOAT8E4M3FNUZ", "FLOAT16"),
    ("FLOAT", "FLOAT8E5M2"),
    ("FLOAT16", "FLOAT8E5M2"),
    ("FLOAT", "FLOAT8E5M2FNUZ"),
    ("FLOAT16", "FLOAT8E5M2FNUZ"),
    ("FLOAT8E5M2", "FLOAT"),
    ("FLOAT8E5M2", "FLOAT16"),
    ("FLOAT8E5M2FNUZ", "FLOAT"),
    ("FLOAT8E5M2FNUZ", "FLOAT16"),
    ("FLOAT", "UINT4"),
    ("FLOAT16", "UINT4"),
    ("FLOAT", "INT4"),
    ("FLOAT16", "INT4"),
    ("UINT4", "FLOAT"),
    ("UINT4", "FLOAT16"),
    ("UINT4", "UINT8"),
    ("INT4", "FLOAT"),
    ("INT4", "FLOAT16"),
    ("INT4", "INT8"),
    ("FLOAT4E2M1", "FLOAT"),
    ("FLOAT4E2M1", "FLOAT16"),
    ("FLOAT", "FLOAT4E2M1"),
    ("FLOAT16", "FLOAT4E2M1"),
    ("FLOAT", "UINT2"),
    ("FLOAT16", "UINT2"),
    ("FLOAT", "INT2"),
    ("FLOAT16", "INT2"),
    ("UINT2", "FLOAT"),
    ("UINT2", "FLOAT16"),
    ("UINT2", "UINT8"),
    ("INT2", "FLOAT"),
    ("INT2", "FLOAT16"),
    ("INT2", "INT8"),
]

f8_types = {"FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ"}

for from_type, to_type in test_cases:
    if from_type == to_type:
        # Skip cases where from_type and to_type are the same
        continue
    from_dtype = getattr(TensorProto, from_type)
    to_dtype = getattr(TensorProto, to_type)
    from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
    to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)

    if from_type == "BFLOAT16" or to_type == "BFLOAT16":
        np_fp32 = np.array(
            [
                "0.47892547",
                "0.48033667",
                "0.49968487",
                "0.81910545",
                "0.47031248",
                "0.816468",
                "0.21087195",
                "0.7229038",
                "NaN",
                "INF",
                "+INF",
                "-INF",
            ],
            dtype=np.float32,
        )
        input_shape = (3, 4)

    elif from_type in f8_types or to_type in f8_types:
        np_fp32 = np.array(
            [
                "0.47892547",
                "0.48033667",
                "0.49968487",
                "0.81910545",
                "0.47031248",
                "0.7229038",
                "1000000",
                "1e-7",
                "NaN",
                "INF",
                "+INF",
                "-INF",
                "-0.0000001",
                "0.0000001",
                "-1000000",
            ],
            dtype=np.float32,
        )
        input_shape = (3, 5)
    elif from_type in ("UINT4", "INT4") or to_type in ("UINT4", "INT4"):
        np_fp32 = np.arange(-9, 16).astype(np.float32)
        input_shape = (5, 5)
    elif from_type in ("UINT2", "INT2") or to_type in ("UINT2", "INT2"):
        np_fp32 = np.arange(-3, 4).astype(np.float32)
        input_shape = (7, 1)
    elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
        np_fp32 = np.array(
            [
                "0.48",
                "0.25",
                "1.05",
                "-3.5",
                "-8",
                "9",
                "1000000",
                "1e-7",
                "NaN",
                "INF",
                "+INF",
                "-INF",
                "-4",
                "0.01",
                "-0.0",
            ],
            dtype=np.float32,
        )
        input_shape = (3, 5)

    else:
        np_fp32 = np.array(
            [
                "0.47892547",
                "0.48033667",
                "0.49968487",
                "0.81910545",
                "0.47031248",
                "0.816468",
                "0.21087195",
                "0.7229038",
                "NaN",
                "INF",
                "+INF",
                "-INF",
            ],
            dtype=np.float32,
        ).reshape([3, 4])
        input_shape = (3, 4)

    if from_type in F8_TYPES:
        np_from = onnx.numpy_helper.saturate_cast(np_fp32, from_np_dtype)
        input = make_tensor(
            "input",
            from_dtype,
            input_shape,
            vals=np_from,
            raw=True,
        )
    elif from_type in FOUR_BIT_TYPES:
        np_from = np_fp32.astype(from_np_dtype)
        packed = onnx.numpy_helper._pack_4bitx2(np_from)
        # No byteswap needed on big-endian machines as _pack_4bitx2()
        # returns a numpy array with uint8 datatype.
        input = make_tensor(
            "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    elif from_type in TWO_BIT_TYPES:
        np_from = np_fp32.astype(from_np_dtype)
        packed = onnx.numpy_helper._pack_2bitx4(np_from)
        # No byteswap needed on big-endian machines as _pack_2bitx4()
        # returns a numpy array with uint8 datatype.
        input = make_tensor(
            "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    else:
        np_from = np_fp32.astype(from_np_dtype)
        input = make_tensor(
            "input", from_dtype, input_shape, vals=np_from, raw=True
        )

    if to_type in F8_TYPES:
        output = make_tensor(
            "output",
            to_dtype,
            input_shape,
            vals=onnx.numpy_helper.saturate_cast(np_from, to_np_dtype),
            raw=True,
        )
    elif to_type in FOUR_BIT_TYPES:
        packed = onnx.numpy_helper._pack_4bitx2(np_from.astype(to_np_dtype))
        # No byteswap needed on big-endian machines as _pack_4bitx2()
        # returns a numpy array with uint8 datatype.
        output = make_tensor(
            "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    elif to_type in TWO_BIT_TYPES:
        packed = onnx.numpy_helper._pack_2bitx4(np_from.astype(to_np_dtype))
        # No byteswap needed on big-endian machines as _pack_2bitx4()
        # returns a numpy array with uint8 datatype.
        output = make_tensor(
            "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
        )
    else:
        output = make_tensor(
            "output",
            to_dtype,
            input_shape,
            vals=np_from.astype(to_np_dtype),
            raw=True,
        )

    like = make_tensor("like", to_dtype, (0,), vals=[])

    node = onnx.helper.make_node(
        "CastLike",
        inputs=["input", "like"],
        outputs=["output"],
    )

    expect(
        node,
        inputs=[input, like],
        outputs=[output],
        name="test_castlike_" + from_type + "_to_" + to_type,
    )

saturate_false

test_cases = itertools.product(
    [
        "FLOAT",
        "FLOAT16",
    ],
    [
        "FLOAT8E4M3FN",
        "FLOAT8E4M3FNUZ",
        "FLOAT8E5M2",
        "FLOAT8E5M2FNUZ",
    ],
)
input_shape = (3, 5)
for from_type, to_type in test_cases:
    from_dtype = getattr(TensorProto, from_type)
    to_dtype = getattr(TensorProto, to_type)
    from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
    to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)
    np_fp32 = np.array(
        [
            "0.47892547",
            "0.48033667",
            "0.49968487",
            "0.81910545",
            "0.47031248",
            "0.7229038",
            "1000000",
            "1e-7",
            "NaN",
            "INF",
            "+INF",
            "-INF",
            "-0.0000001",
            "0.0000001",
            "-1000000",
        ],
        dtype=np.float32,
    )

    input = make_tensor(
        "input",
        from_dtype,
        input_shape,
        vals=np_fp32.astype(from_np_dtype),
        raw=True,
    )
    output = make_tensor(
        "output",
        to_dtype,
        input_shape,
        vals=np_fp32.astype(from_np_dtype).astype(to_np_dtype),
        raw=True,
    )

    like = make_tensor("like", to_dtype, (0,), vals=[])

    node = onnx.helper.make_node(
        "CastLike",
        inputs=["input", "like"],
        outputs=["output"],
        saturate=0,
    )

    expect(
        node,
        inputs=[input, like],
        outputs=[output],
        name="test_castlike_no_saturate_" + from_type + "_to_" + to_type,
    )

CausalConvWithState

There are 13 test cases, listed as following:

b1_c1_degenerate

# Mamba/GDN inner-head edge case: B=1, C=1.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 1, 1, 6, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_b1_c1_degenerate",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

basic

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_basic",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

decode_step

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "bias", "past_state"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 1, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
bias = np.random.randn(channels).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(
    input_, weight, bias=bias, past_state=past_state
)

expect(
    node,
    inputs=[input_, weight, bias, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_decode_step",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

fp16

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.rand(batch_size, channels, length).astype(np.float16)
weight = np.random.rand(channels, 1, k).astype(np.float16)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

kernel_size_one

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 1
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_kernel_size_one",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

short_input_no_past_state

# L < k-1 with no past_state: zero-pad is wider than the input.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 2, 5
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_short_input_no_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

silu

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
    activation="silu",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight, activation="silu")

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_silu",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

silu_fp16

# fp16 + SiLU: the reference upcasts Sigmoid/Mul to float32, so the
# function-body expansion must do the same to stay numerically faithful.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
    activation="silu",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.rand(batch_size, channels, length).astype(np.float16)
weight = np.random.rand(channels, 1, k).astype(np.float16)

output, present_state = _compute(input_, weight, activation="silu")

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_silu_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

silu_with_past_state

# Fused activation combined with concat-from-past variant of PaddedInput.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "", "past_state"],
    outputs=["output", "present_state"],
    activation="silu",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(
    input_, weight, past_state=past_state, activation="silu"
)

expect(
    node,
    inputs=[input_, weight, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_silu_with_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

swish_alias

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
    activation="swish",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight, activation="swish")

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_swish_alias",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

with_bias

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "bias"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
bias = np.random.randn(channels).astype(np.float32)

output, present_state = _compute(input_, weight, bias=bias)

expect(
    node,
    inputs=[input_, weight, bias],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_with_bias",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

with_bias_and_past_state

# Multi-token (T>1) path through Concat(past, input) -> Conv(+bias).
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "bias", "past_state"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
bias = np.random.randn(channels).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(
    input_, weight, bias=bias, past_state=past_state
)

expect(
    node,
    inputs=[input_, weight, bias, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_with_bias_and_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

with_past_state

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "", "past_state"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(input_, weight, past_state=past_state)

expect(
    node,
    inputs=[input_, weight, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_with_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

Ceil

There are 1 test cases, listed as following:

ceil

node = onnx.helper.make_node(
    "Ceil",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1.5, 1.2]).astype(np.float32)
y = np.ceil(x)  # expected output [-1., 2.]
expect(node, inputs=[x], outputs=[y], name="test_ceil_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.ceil(x)
expect(node, inputs=[x], outputs=[y], name="test_ceil")

Celu

There are 3 test cases, listed as following:

celu

alpha = 2.0
node = onnx.helper.make_node(
    "Celu",
    inputs=["X"],
    outputs=["Y"],
    alpha=alpha,
)

input_data = np.array(
    [
        [
            [[0.8439683], [0.5665144], [0.05836735]],
            [[0.02916367], [0.12964272], [0.5060197]],
            [[0.79538304], [0.9411346], [0.9546573]],
        ],
        [
            [[0.17730942], [0.46192095], [0.26480448]],
            [[0.6746842], [0.01665257], [0.62473077]],
            [[0.9240844], [0.9722341], [0.11965699]],
        ],
        [
            [[0.41356155], [0.9129373], [0.59330076]],
            [[0.81929934], [0.7862604], [0.11799799]],
            [[0.69248444], [0.54119414], [0.07513223]],
        ],
    ],
    dtype=np.float32,
)

# Calculate expected output data
positive_input = np.maximum(0, input_data)
negative_input = np.minimum(0, alpha * (np.exp(input_data / alpha) - 1))
expected_output = positive_input + negative_input

expect(node, inputs=[input_data], outputs=[expected_output], name="test_celu")

celu_bfloat16

alpha = 2.0
node = onnx.helper.make_node(
    "Celu",
    inputs=["X"],
    outputs=["Y"],
    alpha=alpha,
)

input_data = np.array([-3.0, -0.5, 0.0, 0.5, 3.0], dtype=ml_dtypes.bfloat16)

positive_input = np.maximum(0, input_data)
negative_input = np.minimum(0, alpha * (np.exp(input_data / alpha) - 1))
expected_output = (positive_input + negative_input).astype(ml_dtypes.bfloat16)

expect(
    node,
    inputs=[input_data],
    outputs=[expected_output],
    name="test_celu_bfloat16",
)

celu_float16

alpha = 2.0
node = onnx.helper.make_node(
    "Celu",
    inputs=["X"],
    outputs=["Y"],
    alpha=alpha,
)

input_data = np.array([-3.0, -0.5, 0.0, 0.5, 3.0], dtype=np.float16)

positive_input = np.maximum(0, input_data)
negative_input = np.minimum(0, alpha * (np.exp(input_data / alpha) - 1))
expected_output = (positive_input + negative_input).astype(np.float16)

expect(
    node,
    inputs=[input_data],
    outputs=[expected_output],
    name="test_celu_float16",
)

CenterCropPad

There are 6 test cases, listed as following:

center_crop_pad_crop

node = onnx.helper.make_node(
    "CenterCropPad",
    inputs=["x", "shape"],
    outputs=["y"],
)

# First dim is even diff, second is uneven
x = np.random.randn(20, 10, 3).astype(np.float32)
shape = np.array([10, 7, 3], dtype=np.int64)
y = x[5:15, 1:8, :]

expect(node, inputs=[x, shape], outputs=[y], name="test_center_crop_pad_crop")

center_crop_pad_crop_and_pad

node = onnx.helper.make_node(
    "CenterCropPad",
    inputs=["x", "shape"],
    outputs=["y"],
)

# Cropping on first dim, padding on second, third stays the same
x = np.random.randn(20, 8, 3).astype(np.float32)
shape = np.array([10, 10, 3], dtype=np.int64)
y = np.zeros([10, 10, 3], dtype=np.float32)
y[:, 1:9, :] = x[5:15, :, :]

expect(
    node,
    inputs=[x, shape],
    outputs=[y],
    name="test_center_crop_pad_crop_and_pad",
)

center_crop_pad_crop_axes_chw

node = onnx.helper.make_node(
    "CenterCropPad",
    inputs=["x", "shape"],
    outputs=["y"],
    axes=[1, 2],
)

# Cropping on second dim, padding on third, first stays the same
x = np.random.randn(3, 20, 8).astype(np.float32)
shape = np.array([10, 9], dtype=np.int64)
y = np.zeros([3, 10, 9], dtype=np.float32)
y[:, :, :8] = x[:, 5:15, :]

expect(
    node,
    inputs=[x, shape],
    outputs=[y],
    name="test_center_crop_pad_crop_axes_chw",
)

center_crop_pad_crop_axes_hwc

node = onnx.helper.make_node(
    "CenterCropPad",
    inputs=["x", "shape"],
    outputs=["y"],
    axes=[0, 1],
)

# Cropping on first dim, padding on second, third stays the same
x = np.random.randn(20, 8, 3).astype(np.float32)
shape = np.array([10, 9], dtype=np.int64)
y = np.zeros([10, 9, 3], dtype=np.float32)
y[:, :8, :] = x[5:15, :, :]

expect(
    node,
    inputs=[x, shape],
    outputs=[y],
    name="test_center_crop_pad_crop_axes_hwc",
)

center_crop_pad_crop_negative_axes_hwc

node = onnx.helper.make_node(
    "CenterCropPad",
    inputs=["x", "shape"],
    outputs=["y"],
    axes=[-3, -2],
)

# Cropping on first dim, padding on second, third stays the same
x = np.random.randn(20, 8, 3).astype(np.float32)
shape = np.array([10, 9], dtype=np.int64)
y = np.zeros([10, 9, 3], dtype=np.float32)
y[:, :8, :] = x[5:15, :, :]

expect(
    node,
    inputs=[x, shape],
    outputs=[y],
    name="test_center_crop_pad_crop_negative_axes_hwc",
)

center_crop_pad_pad

node = onnx.helper.make_node(
    "CenterCropPad",
    inputs=["x", "shape"],
    outputs=["y"],
)

# First dim is even diff, second is uneven
x = np.random.randn(10, 7, 3).astype(np.float32)
shape = np.array([20, 10, 3], dtype=np.int64)
y = np.zeros([20, 10, 3], dtype=np.float32)
y[5:15, 1:8, :] = x

expect(node, inputs=[x, shape], outputs=[y], name="test_center_crop_pad_pad")

Clip

There are 3 test cases, listed as following:

clip

node = onnx.helper.make_node(
    "Clip",
    inputs=["x", "min", "max"],
    outputs=["y"],
)

x = np.array([-2, 0, 2]).astype(np.float32)
min_val = np.float32(-1)
max_val = np.float32(1)
y = np.clip(x, min_val, max_val)  # expected output [-1., 0., 1.]
expect(
    node, inputs=[x, min_val, max_val], outputs=[y], name="test_clip_example"
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, min_val, max_val)
expect(node, inputs=[x, min_val, max_val], outputs=[y], name="test_clip")
node = onnx.helper.make_node(
    "Clip",
    inputs=["x", "min", "max"],
    outputs=["y"],
)

min_val = np.float32(-5)
max_val = np.float32(5)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.array([-1, 0, 1]).astype(np.float32)
expect(
    node, inputs=[x, min_val, max_val], outputs=[y], name="test_clip_inbounds"
)

x = np.array([-6, 0, 6]).astype(np.float32)
y = np.array([-5, 0, 5]).astype(np.float32)
expect(
    node, inputs=[x, min_val, max_val], outputs=[y], name="test_clip_outbounds"
)

x = np.array([-1, 0, 6]).astype(np.float32)
y = np.array([-1, 0, 5]).astype(np.float32)
expect(
    node,
    inputs=[x, min_val, max_val],
    outputs=[y],
    name="test_clip_splitbounds",
)

x = np.array([-2, 0, 6]).astype(np.float32)
y = np.array([1, 1, 1]).astype(np.float32)
min_val = np.float32(2)
max_val = np.float32(1)
expect(
    node,
    inputs=[x, min_val, max_val],
    outputs=[y],
    name="test_clip_min_greater_than_max",
)

clip_default

node = onnx.helper.make_node(
    "Clip",
    inputs=["x", "min"],
    outputs=["y"],
)
min_val = np.float32(0)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, min_val, np.inf)
expect(node, inputs=[x, min_val], outputs=[y], name="test_clip_default_min")

no_min = ""  # optional input, not supplied
node = onnx.helper.make_node(
    "Clip",
    inputs=["x", no_min, "max"],
    outputs=["y"],
)
max_val = np.float32(0)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, -np.inf, max_val)
expect(node, inputs=[x, max_val], outputs=[y], name="test_clip_default_max")

no_max = ""  # optional input, not supplied
node = onnx.helper.make_node(
    "Clip",
    inputs=["x", no_min, no_max],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.array([-1, 0, 1]).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_clip_default_inbounds")

clip_default_int8

node = onnx.helper.make_node(
    "Clip",
    inputs=["x", "min"],
    outputs=["y"],
)
min_val = np.int8(0)
x = np.random.randn(3, 4, 5).astype(np.int8)
y = np.clip(x, min_val, np.iinfo(np.int8).max)
expect(
    node, inputs=[x, min_val], outputs=[y], name="test_clip_default_int8_min"
)

no_min = ""  # optional input, not supplied
node = onnx.helper.make_node(
    "Clip",
    inputs=["x", no_min, "max"],
    outputs=["y"],
)
max_val = np.int8(0)
x = np.random.randn(3, 4, 5).astype(np.int8)
y = np.clip(x, np.iinfo(np.int8).min, max_val)
expect(
    node, inputs=[x, max_val], outputs=[y], name="test_clip_default_int8_max"
)

no_max = ""  # optional input, not supplied
node = onnx.helper.make_node(
    "Clip",
    inputs=["x", no_min, no_max],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.int8)
y = np.array([-1, 0, 1]).astype(np.int8)
expect(node, inputs=[x], outputs=[y], name="test_clip_default_int8_inbounds")

Col2Im

There are 5 test cases, listed as following:

col2im

input = np.array(
    [
        [
            [1.0, 6.0, 11.0, 16.0, 21.0],  # (1, 5, 5)
            [2.0, 7.0, 12.0, 17.0, 22.0],
            [3.0, 8.0, 13.0, 18.0, 23.0],
            [4.0, 9.0, 14.0, 19.0, 24.0],
            [5.0, 0.0, 15.0, 20.0, 25.0],
        ]
    ]
).astype(np.float32)

image_shape = np.array([5, 5]).astype(np.int64)
block_shape = np.array([1, 5]).astype(np.int64)
node = onnx.helper.make_node(
    "Col2Im", ["input", "image_shape", "block_shape"], ["output"]
)

output = np.array(
    [
        [
            [
                [1.0, 2.0, 3.0, 4.0, 5.0],  # (1, 1, 5, 5)
                [6.0, 7.0, 8.0, 9.0, 0.0],
                [11.0, 12.0, 13.0, 14.0, 15.0],
                [16.0, 17.0, 18.0, 19.0, 20.0],
                [21.0, 22.0, 23.0, 24.0, 25.0],
            ]
        ]
    ]
).astype(np.float32)

expect(
    node,
    inputs=[input, image_shape, block_shape],
    outputs=[output],
    name="test_col2im",
)

col2im_5d

input = np.array(
    [
        [
            [1, 6, 11, 16, 21, 26, 31, 36, 41, 46, 51, 56],  # (1, 10, 12)
            [2, 7, 12, 17, 22, 27, 32, 37, 42, 47, 52, 57],
            [3, 8, 13, 18, 23, 28, 33, 38, 43, 48, 53, 58],
            [4, 9, 14, 19, 24, 29, 34, 39, 44, 49, 54, 59],
            [5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60],
            [61, 66, 71, 76, 81, 86, 91, 96, 101, 106, 111, 116],
            [62, 67, 72, 77, 82, 87, 92, 97, 102, 107, 112, 117],
            [63, 68, 73, 78, 83, 88, 93, 98, 103, 108, 113, 118],
            [64, 69, 74, 79, 84, 89, 94, 99, 104, 109, 114, 119],
            [65, 70, 75, 80, 85, 90, 95, 100, 105, 110, 115, 120],
        ]
    ]
).astype(np.float32)
image_shape = np.array([3, 4, 5]).astype(np.int64)
block_shape = np.array([1, 1, 5]).astype(np.int64)

output = np.array(
    [
        [
            [
                [
                    [1, 2, 3, 4, 5],  # (1, 2, 3, 4, 5)
                    [6, 7, 8, 9, 10],
                    [11, 12, 13, 14, 15],
                    [16, 17, 18, 19, 20],
                ],
                [
                    [21, 22, 23, 24, 25],
                    [26, 27, 28, 29, 30],
                    [31, 32, 33, 34, 35],
                    [36, 37, 38, 39, 40],
                ],
                [
                    [41, 42, 43, 44, 45],
                    [46, 47, 48, 49, 50],
                    [51, 52, 53, 54, 55],
                    [56, 57, 58, 59, 60],
                ],
            ],
            [
                [
                    [61, 62, 63, 64, 65],
                    [66, 67, 68, 69, 70],
                    [71, 72, 73, 74, 75],
                    [76, 77, 78, 79, 80],
                ],
                [
                    [81, 82, 83, 84, 85],
                    [86, 87, 88, 89, 90],
                    [91, 92, 93, 94, 95],
                    [96, 97, 98, 99, 100],
                ],
                [
                    [101, 102, 103, 104, 105],
                    [106, 107, 108, 109, 110],
                    [111, 112, 113, 114, 115],
                    [116, 117, 118, 119, 120],
                ],
            ],
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "Col2Im", ["input", "image_shape", "block_shape"], ["output"]
)
expect(
    node,
    inputs=[input, image_shape, block_shape],
    outputs=[output],
    name="test_col2im_5d",
)

col2im_dilations

input = np.array(
    [
        [
            [1.0, 5.0, 9.0, 13.0, 17],  # (1, 4, 5)
            [2.0, 6.0, 10.0, 14.0, 18],
            [3.0, 7.0, 11.0, 15.0, 19],
            [4.0, 8.0, 12.0, 16.0, 20],
        ]
    ]
).astype(np.float32)
image_shape = np.array([6, 6]).astype(np.int64)
block_shape = np.array([2, 2]).astype(np.int64)

output = np.array(
    [
        [
            [
                [1.0, 0.0, 0.0, 0.0, 0.0, 2.0],  # (1, 1, 6, 6)
                [8.0, 0.0, 0.0, 0.0, 0.0, 10.0],
                [16.0, 0.0, 0.0, 0.0, 0.0, 18.0],
                [24.0, 0.0, 0.0, 0.0, 0.0, 26.0],
                [32.0, 0.0, 0.0, 0.0, 0.0, 34.0],
                [19.0, 0.0, 0.0, 0.0, 0.0, 20.0],
            ]
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "Col2Im",
    ["input", "image_shape", "block_shape"],
    ["output"],
    dilations=[1, 5],
)
expect(
    node,
    inputs=[input, image_shape, block_shape],
    outputs=[output],
    name="test_col2im_dilations",
)

col2im_pads

input = np.array(
    [
        [
            [
                1.0,
                6.0,
                11.0,
                16.0,
                21.0,
                26,
                31,
                36,
                41,
                46,
                51,
                56,
                61,
                66,
                71,
            ],  # (1, 5, 15)
            [
                2.0,
                7.0,
                12.0,
                17.0,
                22.0,
                27,
                32,
                37,
                42,
                47,
                52,
                57,
                62,
                67,
                72,
            ],
            [
                3.0,
                8.0,
                13.0,
                18.0,
                23.0,
                28,
                33,
                38,
                43,
                48,
                53,
                58,
                63,
                68,
                73,
            ],
            [
                4.0,
                9.0,
                14.0,
                19.0,
                24.0,
                29,
                34,
                39,
                44,
                49,
                54,
                59,
                64,
                69,
                74,
            ],
            [
                5.0,
                10.0,
                15.0,
                20.0,
                25.0,
                30,
                35,
                40,
                45,
                50,
                55,
                60,
                65,
                70,
                75,
            ],
        ]
    ]
).astype(np.float32)
image_shape = np.array([5, 5]).astype(np.int64)
block_shape = np.array([1, 5]).astype(np.int64)

output = np.array(
    [
        [
            [
                [8.0, 21.0, 24.0, 27.0, 24.0],  # (1, 1, 5, 5)
                [38.0, 66.0, 69.0, 72.0, 54.0],
                [68.0, 111.0, 114.0, 117.0, 84.0],
                [98.0, 156.0, 159.0, 162.0, 114.0],
                [128.0, 201.0, 204.0, 207.0, 144.0],
            ]
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "Col2Im",
    ["input", "image_shape", "block_shape"],
    ["output"],
    pads=[0, 1, 0, 1],
)
expect(
    node,
    inputs=[input, image_shape, block_shape],
    outputs=[output],
    name="test_col2im_pads",
)

col2im_strides

input = np.array(
    [
        [
            [0.0, 0.0, 0.0, 0.0],  # (1, 9, 4)
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0, 0.0],
        ]
    ]
).astype(np.float32)
image_shape = np.array([5, 5]).astype(np.int64)
block_shape = np.array([3, 3]).astype(np.int64)

output = np.array(
    [
        [
            [
                [0.0, 1.0, 1.0, 1.0, 1.0],  # (1, 1, 5, 5)
                [1.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 2.0, 1.0, 2.0, 1.0],
                [1.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 1.0, 0.0],
            ]
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "Col2Im",
    ["input", "image_shape", "block_shape"],
    ["output"],
    strides=[2, 2],
)
expect(
    node,
    inputs=[input, image_shape, block_shape],
    outputs=[output],
    name="test_col2im_strides",
)

Compress

There are 4 test cases, listed as following:

compress_0

node = onnx.helper.make_node(
    "Compress",
    inputs=["input", "condition"],
    outputs=["output"],
    axis=0,
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1, 1])
output = np.compress(condition, input, axis=0)
# print(output)
# [[ 3.  4.]
# [ 5.  6.]]

expect(
    node,
    inputs=[input, condition.astype(bool)],
    outputs=[output],
    name="test_compress_0",
)

compress_1

node = onnx.helper.make_node(
    "Compress",
    inputs=["input", "condition"],
    outputs=["output"],
    axis=1,
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1])
output = np.compress(condition, input, axis=1)
# print(output)
# [[ 2.]
# [ 4.]
# [ 6.]]

expect(
    node,
    inputs=[input, condition.astype(bool)],
    outputs=[output],
    name="test_compress_1",
)

compress_default_axis

node = onnx.helper.make_node(
    "Compress",
    inputs=["input", "condition"],
    outputs=["output"],
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1, 0, 0, 1])
output = np.compress(condition, input)
# print(output)
# [ 2., 5.]

expect(
    node,
    inputs=[input, condition.astype(bool)],
    outputs=[output],
    name="test_compress_default_axis",
)

compress_negative_axis

node = onnx.helper.make_node(
    "Compress",
    inputs=["input", "condition"],
    outputs=["output"],
    axis=-1,
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1])
output = np.compress(condition, input, axis=-1)
# print(output)
# [[ 2.]
# [ 4.]
# [ 6.]]
expect(
    node,
    inputs=[input, condition.astype(bool)],
    outputs=[output],
    name="test_compress_negative_axis",
)

Concat

There are 1 test cases, listed as following:

concat

test_cases: dict[str, Sequence[Any]] = {
    "1d": ([1, 2], [3, 4]),
    "2d": ([[1, 2], [3, 4]], [[5, 6], [7, 8]]),
    "3d": (
        [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
        [[[9, 10], [11, 12]], [[13, 14], [15, 16]]],
    ),
}

for test_case, values_ in test_cases.items():
    values = [np.asarray(v, dtype=np.float32) for v in values_]
    for i in range(len(values[0].shape)):
        in_args = ["value" + str(k) for k in range(len(values))]
        node = onnx.helper.make_node(
            "Concat", inputs=list(in_args), outputs=["output"], axis=i
        )
        output = np.concatenate(values, i)
        expect(
            node,
            inputs=list(values),
            outputs=[output],
            name="test_concat_" + test_case + "_axis_" + str(i),
        )

    for i in range(-len(values[0].shape), 0):
        in_args = ["value" + str(k) for k in range(len(values))]
        node = onnx.helper.make_node(
            "Concat", inputs=list(in_args), outputs=["output"], axis=i
        )
        output = np.concatenate(values, i)
        expect(
            node,
            inputs=list(values),
            outputs=[output],
            name="test_concat_" + test_case + "_axis_negative_" + str(abs(i)),
        )

Constant

There are 1 test cases, listed as following:

constant

values = np.random.randn(5, 5).astype(np.float32)
node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["values"],
    value=onnx.helper.make_tensor(
        name="const_tensor",
        data_type=onnx.TensorProto.FLOAT,
        dims=values.shape,
        vals=values.flatten().astype(float),
    ),
)

expect(node, inputs=[], outputs=[values], name="test_constant")

ConstantOfShape

There are 3 test cases, listed as following:

float_ones

x = np.array([4, 3, 2]).astype(np.int64)
tensor_value = onnx.helper.make_tensor(
    "value", onnx.TensorProto.FLOAT, [1], [1]
)
node = onnx.helper.make_node(
    "ConstantOfShape",
    inputs=["x"],
    outputs=["y"],
    value=tensor_value,
)

y = np.ones(x, dtype=np.float32)
expect(node, inputs=[x], outputs=[y], name="test_constantofshape_float_ones")

int32_shape_zero

x = np.array(
    [
        0,
    ]
).astype(np.int64)
tensor_value = onnx.helper.make_tensor(
    "value", onnx.TensorProto.INT32, [1], [0]
)
node = onnx.helper.make_node(
    "ConstantOfShape",
    inputs=["x"],
    outputs=["y"],
    value=tensor_value,
)
y = np.zeros(x, dtype=np.int32)
expect(
    node, inputs=[x], outputs=[y], name="test_constantofshape_int_shape_zero"
)

int32_zeros

x = np.array([10, 6]).astype(np.int64)
tensor_value = onnx.helper.make_tensor(
    "value", onnx.TensorProto.INT32, [1], [0]
)
node = onnx.helper.make_node(
    "ConstantOfShape",
    inputs=["x"],
    outputs=["y"],
    value=tensor_value,
)
y = np.zeros(x, dtype=np.int32)
expect(node, inputs=[x], outputs=[y], name="test_constantofshape_int_zeros")

Conv

There are 3 test cases, listed as following:

conv

x = np.array(
    [
        [
            [
                [0.0, 1.0, 2.0, 3.0, 4.0],  # (1, 1, 5, 5) input tensor
                [5.0, 6.0, 7.0, 8.0, 9.0],
                [10.0, 11.0, 12.0, 13.0, 14.0],
                [15.0, 16.0, 17.0, 18.0, 19.0],
                [20.0, 21.0, 22.0, 23.0, 24.0],
            ]
        ]
    ]
).astype(np.float32)
W = np.array(
    [
        [
            [
                [1.0, 1.0, 1.0],  # (1, 1, 3, 3) tensor for convolution weights
                [1.0, 1.0, 1.0],
                [1.0, 1.0, 1.0],
            ]
        ]
    ]
).astype(np.float32)

# Convolution with padding
node_with_padding = onnx.helper.make_node(
    "Conv",
    inputs=["x", "W"],
    outputs=["y"],
    kernel_shape=[3, 3],
    # Default values for other attributes: strides=[1, 1], dilations=[1, 1], groups=1
    pads=[1, 1, 1, 1],
)
y_with_padding = np.array(
    [
        [
            [
                [12.0, 21.0, 27.0, 33.0, 24.0],  # (1, 1, 5, 5) output tensor
                [33.0, 54.0, 63.0, 72.0, 51.0],
                [63.0, 99.0, 108.0, 117.0, 81.0],
                [93.0, 144.0, 153.0, 162.0, 111.0],
                [72.0, 111.0, 117.0, 123.0, 84.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_with_padding,
    inputs=[x, W],
    outputs=[y_with_padding],
    name="test_basic_conv_with_padding",
)

# Convolution without padding
node_without_padding = onnx.helper.make_node(
    "Conv",
    inputs=["x", "W"],
    outputs=["y"],
    kernel_shape=[3, 3],
    # Default values for other attributes: strides=[1, 1], dilations=[1, 1], groups=1
    pads=[0, 0, 0, 0],
)
y_without_padding = np.array(
    [
        [
            [
                [54.0, 63.0, 72.0],  # (1, 1, 3, 3) output tensor
                [99.0, 108.0, 117.0],
                [144.0, 153.0, 162.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_without_padding,
    inputs=[x, W],
    outputs=[y_without_padding],
    name="test_basic_conv_without_padding",
)

conv_with_autopad_same

x = np.array(
    [
        [
            [
                [0.0, 1.0, 2.0, 3.0, 4.0],  # (1, 1, 5, 5) input tensor
                [5.0, 6.0, 7.0, 8.0, 9.0],
                [10.0, 11.0, 12.0, 13.0, 14.0],
                [15.0, 16.0, 17.0, 18.0, 19.0],
                [20.0, 21.0, 22.0, 23.0, 24.0],
            ]
        ]
    ]
).astype(np.float32)
W = np.array(
    [
        [
            [
                [1.0, 1.0, 1.0],  # (1, 1, 3, 3) tensor for convolution weights
                [1.0, 1.0, 1.0],
                [1.0, 1.0, 1.0],
            ]
        ]
    ]
).astype(np.float32)

# Convolution with auto_pad='SAME_LOWER' and strides=2
node = onnx.helper.make_node(
    "Conv",
    inputs=["x", "W"],
    outputs=["y"],
    auto_pad="SAME_LOWER",
    kernel_shape=[3, 3],
    strides=[2, 2],
)
y = np.array(
    [[[[12.0, 27.0, 24.0], [63.0, 108.0, 81.0], [72.0, 117.0, 84.0]]]]
).astype(np.float32)
expect(node, inputs=[x, W], outputs=[y], name="test_conv_with_autopad_same")

conv_with_strides

x = np.array(
    [
        [
            [
                [0.0, 1.0, 2.0, 3.0, 4.0],  # (1, 1, 7, 5) input tensor
                [5.0, 6.0, 7.0, 8.0, 9.0],
                [10.0, 11.0, 12.0, 13.0, 14.0],
                [15.0, 16.0, 17.0, 18.0, 19.0],
                [20.0, 21.0, 22.0, 23.0, 24.0],
                [25.0, 26.0, 27.0, 28.0, 29.0],
                [30.0, 31.0, 32.0, 33.0, 34.0],
            ]
        ]
    ]
).astype(np.float32)
W = np.array(
    [
        [
            [
                [1.0, 1.0, 1.0],  # (1, 1, 3, 3) tensor for convolution weights
                [1.0, 1.0, 1.0],
                [1.0, 1.0, 1.0],
            ]
        ]
    ]
).astype(np.float32)

# Convolution with strides=2 and padding
node_with_padding = onnx.helper.make_node(
    "Conv",
    inputs=["x", "W"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[1, 1, 1, 1],
    strides=[
        2,
        2,
    ],  # Default values for other attributes: dilations=[1, 1], groups=1
)
y_with_padding = np.array(
    [
        [
            [
                [12.0, 27.0, 24.0],  # (1, 1, 4, 3) output tensor
                [63.0, 108.0, 81.0],
                [123.0, 198.0, 141.0],
                [112.0, 177.0, 124.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_with_padding,
    inputs=[x, W],
    outputs=[y_with_padding],
    name="test_conv_with_strides_padding",
)

# Convolution with strides=2 and no padding
node_without_padding = onnx.helper.make_node(
    "Conv",
    inputs=["x", "W"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[0, 0, 0, 0],
    strides=[
        2,
        2,
    ],  # Default values for other attributes: dilations=[1, 1], groups=1
)
y_without_padding = np.array(
    [
        [
            [
                [54.0, 72.0],  # (1, 1, 3, 2) output tensor
                [144.0, 162.0],
                [234.0, 252.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_without_padding,
    inputs=[x, W],
    outputs=[y_without_padding],
    name="test_conv_with_strides_no_padding",
)

# Convolution with strides=2 and padding only along one dimension (the H dimension in NxCxHxW tensor)
node_with_asymmetric_padding = onnx.helper.make_node(
    "Conv",
    inputs=["x", "W"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[1, 0, 1, 0],
    strides=[
        2,
        2,
    ],  # Default values for other attributes: dilations=[1, 1], groups=1
)
y_with_asymmetric_padding = np.array(
    [
        [
            [
                [21.0, 33.0],  # (1, 1, 4, 2) output tensor
                [99.0, 117.0],
                [189.0, 207.0],
                [171.0, 183.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_with_asymmetric_padding,
    inputs=[x, W],
    outputs=[y_with_asymmetric_padding],
    name="test_conv_with_strides_and_asymmetric_padding",
)

ConvInteger

There are 2 test cases, listed as following:

with_padding

x = (
    np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
    .astype(np.uint8)
    .reshape((1, 1, 3, 3))
)
x_zero_point = np.uint8(1)
w_zero_points = np.array([0, 1], dtype=np.uint8)
w = np.array([1, 1, 1, 1, 1, 1, 1, 1]).astype(np.uint8).reshape((2, 1, 2, 2))

y = (
    np.array(
        [
            1,
            3,
            5,
            3,
            5,
            12,
            16,
            9,
            11,
            24,
            28,
            15,
            7,
            15,
            17,
            9,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
        ]
    )
    .astype(np.int32)
    .reshape((1, 2, 4, 4))
)

# ConvInteger with padding
convinteger_node_with_padding = onnx.helper.make_node(
    "ConvInteger",
    inputs=["x", "w", "x_zero_point", "w_zero_points"],
    outputs=["y"],
    pads=[1, 1, 1, 1],
)

expect(
    convinteger_node_with_padding,
    inputs=[x, w, x_zero_point, w_zero_points],
    outputs=[y],
    name="test_convinteger_with_padding",
)

without_padding

x = (
    np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
    .astype(np.uint8)
    .reshape((1, 1, 3, 3))
)
x_zero_point = np.uint8(1)
w = np.array([1, 1, 1, 1]).astype(np.uint8).reshape((1, 1, 2, 2))

y = np.array([12, 16, 24, 28]).astype(np.int32).reshape(1, 1, 2, 2)

# ConvInteger without padding
convinteger_node = onnx.helper.make_node(
    "ConvInteger", inputs=["x", "w", "x_zero_point"], outputs=["y"]
)

expect(
    convinteger_node,
    inputs=[x, w, x_zero_point],
    outputs=[y],
    name="test_convinteger_without_padding",
)

ConvTranspose

There are 9 test cases, listed as following:

convtranspose

x = np.array(
    [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]]  # (1, 1, 3, 3)
).astype(np.float32)

W = np.array(
    [
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],  # (1, 2, 3, 3)
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node("ConvTranspose", ["X", "W"], ["Y"])

y = np.array(
    [
        [
            [
                [0.0, 1.0, 3.0, 3.0, 2.0],  # (1, 2, 5, 5)
                [3.0, 8.0, 15.0, 12.0, 7.0],
                [9.0, 21.0, 36.0, 27.0, 15.0],
                [9.0, 20.0, 33.0, 24.0, 13.0],
                [6.0, 13.0, 21.0, 15.0, 8.0],
            ],
            [
                [0.0, 1.0, 3.0, 3.0, 2.0],
                [3.0, 8.0, 15.0, 12.0, 7.0],
                [9.0, 21.0, 36.0, 27.0, 15.0],
                [9.0, 20.0, 33.0, 24.0, 13.0],
                [6.0, 13.0, 21.0, 15.0, 8.0],
            ],
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose")

convtranspose_1d

x = np.array([[[0.0, 1.0, 2.0]]]).astype(np.float32)  # (1, 1, 3)

W = np.array([[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]]).astype(  # (1, 2, 3)
    np.float32
)

node = onnx.helper.make_node("ConvTranspose", ["X", "W"], ["Y"])

y = np.array(
    [[[0.0, 1.0, 3.0, 3.0, 2.0], [0.0, 1.0, 3.0, 3.0, 2.0]]]  # (1, 2, 5)
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_1d")

convtranspose_3d

x = np.array(
    [
        [
            [
                [
                    [0.0, 1.0, 2.0, 3.0, 4.0],  # (1, 1, 3, 4, 5)
                    [5.0, 6.0, 7.0, 8.0, 9.0],
                    [10.0, 11.0, 12.0, 13.0, 14.0],
                    [15.0, 16.0, 17.0, 18.0, 19.0],
                ],
                [
                    [20.0, 21.0, 22.0, 23.0, 24.0],
                    [25.0, 26.0, 27.0, 28.0, 29.0],
                    [30.0, 31.0, 32.0, 33.0, 34.0],
                    [35.0, 36.0, 37.0, 38.0, 39.0],
                ],
                [
                    [40.0, 41.0, 42.0, 43.0, 44.0],
                    [45.0, 46.0, 47.0, 48.0, 49.0],
                    [50.0, 51.0, 52.0, 53.0, 54.0],
                    [55.0, 56.0, 57.0, 58.0, 59.0],
                ],
            ]
        ]
    ]
).astype(np.float32)

W = np.array(
    [
        [
            [
                [
                    [1.0, 1.0, 1.0],  # (1, 2, 3, 3, 3)
                    [1.0, 1.0, 1.0],
                    [1.0, 1.0, 1.0],
                ],
                [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
                [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
            ],
            [
                [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
                [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
                [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
            ],
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node("ConvTranspose", ["X", "W"], ["Y"])

y = np.array(
    [
        [
            [
                [
                    [0.0, 1.0, 3.0, 6.0, 9.0, 7.0, 4.0],  # (1, 2, 5, 6, 7)
                    [5.0, 12.0, 21.0, 27.0, 33.0, 24.0, 13.0],
                    [15.0, 33.0, 54.0, 63.0, 72.0, 51.0, 27.0],
                    [30.0, 63.0, 99.0, 108.0, 117.0, 81.0, 42.0],
                    [25.0, 52.0, 81.0, 87.0, 93.0, 64.0, 33.0],
                    [15.0, 31.0, 48.0, 51.0, 54.0, 37.0, 19.0],
                ],
                [
                    [20.0, 42.0, 66.0, 72.0, 78.0, 54.0, 28.0],
                    [50.0, 104.0, 162.0, 174.0, 186.0, 128.0, 66.0],
                    [90.0, 186.0, 288.0, 306.0, 324.0, 222.0, 114.0],
                    [120.0, 246.0, 378.0, 396.0, 414.0, 282.0, 144.0],
                    [90.0, 184.0, 282.0, 294.0, 306.0, 208.0, 106.0],
                    [50.0, 102.0, 156.0, 162.0, 168.0, 114.0, 58.0],
                ],
                [
                    [60.0, 123.0, 189.0, 198.0, 207.0, 141.0, 72.0],
                    [135.0, 276.0, 423.0, 441.0, 459.0, 312.0, 159.0],
                    [225.0, 459.0, 702.0, 729.0, 756.0, 513.0, 261.0],
                    [270.0, 549.0, 837.0, 864.0, 891.0, 603.0, 306.0],
                    [195.0, 396.0, 603.0, 621.0, 639.0, 432.0, 219.0],
                    [105.0, 213.0, 324.0, 333.0, 342.0, 231.0, 117.0],
                ],
                [
                    [60.0, 122.0, 186.0, 192.0, 198.0, 134.0, 68.0],
                    [130.0, 264.0, 402.0, 414.0, 426.0, 288.0, 146.0],
                    [210.0, 426.0, 648.0, 666.0, 684.0, 462.0, 234.0],
                    [240.0, 486.0, 738.0, 756.0, 774.0, 522.0, 264.0],
                    [170.0, 344.0, 522.0, 534.0, 546.0, 368.0, 186.0],
                    [90.0, 182.0, 276.0, 282.0, 288.0, 194.0, 98.0],
                ],
                [
                    [40.0, 81.0, 123.0, 126.0, 129.0, 87.0, 44.0],
                    [85.0, 172.0, 261.0, 267.0, 273.0, 184.0, 93.0],
                    [135.0, 273.0, 414.0, 423.0, 432.0, 291.0, 147.0],
                    [150.0, 303.0, 459.0, 468.0, 477.0, 321.0, 162.0],
                    [105.0, 212.0, 321.0, 327.0, 333.0, 224.0, 113.0],
                    [55.0, 111.0, 168.0, 171.0, 174.0, 117.0, 59.0],
                ],
            ],
            [
                [
                    [0.0, 1.0, 3.0, 6.0, 9.0, 7.0, 4.0],
                    [5.0, 12.0, 21.0, 27.0, 33.0, 24.0, 13.0],
                    [15.0, 33.0, 54.0, 63.0, 72.0, 51.0, 27.0],
                    [30.0, 63.0, 99.0, 108.0, 117.0, 81.0, 42.0],
                    [25.0, 52.0, 81.0, 87.0, 93.0, 64.0, 33.0],
                    [15.0, 31.0, 48.0, 51.0, 54.0, 37.0, 19.0],
                ],
                [
                    [20.0, 42.0, 66.0, 72.0, 78.0, 54.0, 28.0],
                    [50.0, 104.0, 162.0, 174.0, 186.0, 128.0, 66.0],
                    [90.0, 186.0, 288.0, 306.0, 324.0, 222.0, 114.0],
                    [120.0, 246.0, 378.0, 396.0, 414.0, 282.0, 144.0],
                    [90.0, 184.0, 282.0, 294.0, 306.0, 208.0, 106.0],
                    [50.0, 102.0, 156.0, 162.0, 168.0, 114.0, 58.0],
                ],
                [
                    [60.0, 123.0, 189.0, 198.0, 207.0, 141.0, 72.0],
                    [135.0, 276.0, 423.0, 441.0, 459.0, 312.0, 159.0],
                    [225.0, 459.0, 702.0, 729.0, 756.0, 513.0, 261.0],
                    [270.0, 549.0, 837.0, 864.0, 891.0, 603.0, 306.0],
                    [195.0, 396.0, 603.0, 621.0, 639.0, 432.0, 219.0],
                    [105.0, 213.0, 324.0, 333.0, 342.0, 231.0, 117.0],
                ],
                [
                    [60.0, 122.0, 186.0, 192.0, 198.0, 134.0, 68.0],
                    [130.0, 264.0, 402.0, 414.0, 426.0, 288.0, 146.0],
                    [210.0, 426.0, 648.0, 666.0, 684.0, 462.0, 234.0],
                    [240.0, 486.0, 738.0, 756.0, 774.0, 522.0, 264.0],
                    [170.0, 344.0, 522.0, 534.0, 546.0, 368.0, 186.0],
                    [90.0, 182.0, 276.0, 282.0, 288.0, 194.0, 98.0],
                ],
                [
                    [40.0, 81.0, 123.0, 126.0, 129.0, 87.0, 44.0],
                    [85.0, 172.0, 261.0, 267.0, 273.0, 184.0, 93.0],
                    [135.0, 273.0, 414.0, 423.0, 432.0, 291.0, 147.0],
                    [150.0, 303.0, 459.0, 468.0, 477.0, 321.0, 162.0],
                    [105.0, 212.0, 321.0, 327.0, 333.0, 224.0, 113.0],
                    [55.0, 111.0, 168.0, 171.0, 174.0, 117.0, 59.0],
                ],
            ],
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_3d")

convtranspose_attributes

x = np.array(
    [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]]  # (1, 1, 3, 3)
).astype(np.float32)

W = np.array(
    [
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],  # (1, 2, 3, 3)
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ]
    ]
).astype(np.float32)

y = np.array(
    [
        [
            [
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0, 2.0, 0.0],  # (1, 2, 10, 8)
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0, 2.0, 0.0],
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0, 2.0, 0.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0, 5.0, 0.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0, 5.0, 0.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0, 5.0, 0.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0, 8.0, 0.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0, 8.0, 0.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0, 8.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            ],
            [
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0, 2.0, 0.0],
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0, 2.0, 0.0],
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0, 2.0, 0.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0, 5.0, 0.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0, 5.0, 0.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0, 5.0, 0.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0, 8.0, 0.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0, 8.0, 0.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0, 8.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            ],
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "ConvTranspose", ["X", "W"], ["Y"], strides=[3, 2], output_shape=[10, 8]
)
expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_output_shape")

node = onnx.helper.make_node(
    "ConvTranspose", ["X", "W"], ["Y"], strides=[3, 2], output_padding=[1, 1]
)
expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_pad")

node = onnx.helper.make_node(
    "ConvTranspose",
    ["X", "W"],
    ["Y"],
    name="test",
    strides=[3, 2],
    output_shape=[10, 8],
    kernel_shape=[3, 3],
    output_padding=[1, 1],
)
expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_kernel_shape")

convtranspose_autopad_same

x = np.array(
    [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]]  # (1, 1, 3, 3)
).astype(np.float32)

W = np.array(
    [
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],  # (1, 2, 3, 3)
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "ConvTranspose", ["X", "W"], ["Y"], auto_pad="SAME_UPPER", strides=[2, 2]
)

y = np.array(
    [
        [
            [
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0],
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0],
                [3.0, 3.0, 8.0, 5.0, 12.0, 7.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0],
                [9.0, 9.0, 20.0, 11.0, 24.0, 13.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0],
            ],
            [
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0],
                [0.0, 0.0, 1.0, 1.0, 3.0, 2.0],
                [3.0, 3.0, 8.0, 5.0, 12.0, 7.0],
                [3.0, 3.0, 7.0, 4.0, 9.0, 5.0],
                [9.0, 9.0, 20.0, 11.0, 24.0, 13.0],
                [6.0, 6.0, 13.0, 7.0, 15.0, 8.0],
            ],
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_autopad_same")

convtranspose_dilations

x = np.array(
    [[[[3.0, 8.0, 1.0], [9.0, 5.0, 7.0], [3.0, 2.0, 6.0]]]]  # (1, 1, 3, 3)
).astype(np.float32)
W = np.array([[[[7.0, 2.0], [1.0, 9.0]]]]).astype(np.float32)  # (1, 1, 2, 2)

node = onnx.helper.make_node(
    "ConvTranspose", ["X", "W"], ["Y"], dilations=[2, 2]
)

y = np.array(
    [
        [
            [
                [21.0, 56.0, 13.0, 16.0, 2.0],  # [1, 1, 5, 5]
                [63.0, 35.0, 67.0, 10.0, 14.0],
                [24.0, 22.0, 76.0, 76.0, 21.0],
                [9.0, 5.0, 88.0, 45.0, 63.0],
                [3.0, 2.0, 33.0, 18.0, 54.0],
            ]
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_dilations")

convtranspose_group_2

x = np.array(
    [
        [
            [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
            [[9.0, 10.0, 11.0], [12.0, 13.0, 14.0], [15.0, 16.0, 17.0]],
        ]
    ]
).astype(np.float32)
W = np.array(
    [
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ],
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ],
    ]
).astype(np.float32)

node = onnx.helper.make_node("ConvTranspose", ["X", "W"], ["Y"], group=2)

y = np.array(
    [
        [
            [
                [0.0, 1.0, 3.0, 3.0, 2.0],
                [3.0, 8.0, 15.0, 12.0, 7.0],
                [9.0, 21.0, 36.0, 27.0, 15.0],
                [9.0, 20.0, 33.0, 24.0, 13.0],
                [6.0, 13.0, 21.0, 15.0, 8.0],
            ],
            [
                [9.0, 19.0, 30.0, 21.0, 11.0],
                [21.0, 44.0, 69.0, 48.0, 25.0],
                [36.0, 75.0, 117.0, 81.0, 42.0],
                [27.0, 56.0, 87.0, 60.0, 31.0],
                [15.0, 31.0, 48.0, 33.0, 17.0],
            ],
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_group_2")

convtranspose_group_2_image_3

x = np.array(
    [
        [
            [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
            [[9.0, 10.0, 11.0], [12.0, 13.0, 14.0], [15.0, 16.0, 17.0]],
        ],
        [
            [[18.0, 19.0, 20.0], [21.0, 22.0, 23.0], [24.0, 25.0, 26.0]],
            [[9.0, 10.0, 11.0], [12.0, 13.0, 14.0], [15.0, 16.0, 17.0]],
        ],
        [
            [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
            [[9.0, 10.0, 11.0], [12.0, 13.0, 14.0], [15.0, 16.0, 17.0]],
        ],
    ]
).astype(np.float32)
W = np.array(
    [
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ],
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ],
    ]
).astype(np.float32)

node = onnx.helper.make_node("ConvTranspose", ["X", "W"], ["Y"], group=2)

y = np.array(
    [
        [
            [
                [0.0, 1.0, 3.0, 3.0, 2.0],
                [3.0, 8.0, 15.0, 12.0, 7.0],
                [9.0, 21.0, 36.0, 27.0, 15.0],
                [9.0, 20.0, 33.0, 24.0, 13.0],
                [6.0, 13.0, 21.0, 15.0, 8.0],
            ],
            [
                [9.0, 19.0, 30.0, 21.0, 11.0],
                [21.0, 44.0, 69.0, 48.0, 25.0],
                [36.0, 75.0, 117.0, 81.0, 42.0],
                [27.0, 56.0, 87.0, 60.0, 31.0],
                [15.0, 31.0, 48.0, 33.0, 17.0],
            ],
        ],
        [
            [
                [18.0, 37.0, 57.0, 39.0, 20.0],
                [39.0, 80.0, 123.0, 84.0, 43.0],
                [63.0, 129.0, 198.0, 135.0, 69.0],
                [45.0, 92.0, 141.0, 96.0, 49.0],
                [24.0, 49.0, 75.0, 51.0, 26.0],
            ],
            [
                [9.0, 19.0, 30.0, 21.0, 11.0],
                [21.0, 44.0, 69.0, 48.0, 25.0],
                [36.0, 75.0, 117.0, 81.0, 42.0],
                [27.0, 56.0, 87.0, 60.0, 31.0],
                [15.0, 31.0, 48.0, 33.0, 17.0],
            ],
        ],
        [
            [
                [0.0, 1.0, 3.0, 3.0, 2.0],
                [3.0, 8.0, 15.0, 12.0, 7.0],
                [9.0, 21.0, 36.0, 27.0, 15.0],
                [9.0, 20.0, 33.0, 24.0, 13.0],
                [6.0, 13.0, 21.0, 15.0, 8.0],
            ],
            [
                [9.0, 19.0, 30.0, 21.0, 11.0],
                [21.0, 44.0, 69.0, 48.0, 25.0],
                [36.0, 75.0, 117.0, 81.0, 42.0],
                [27.0, 56.0, 87.0, 60.0, 31.0],
                [15.0, 31.0, 48.0, 33.0, 17.0],
            ],
        ],
    ]
).astype(np.float32)

expect(
    node, inputs=[x, W], outputs=[y], name="test_convtranspose_group_2_image_3"
)

convtranspose_pads

x = np.array(
    [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]]  # (1, 1, 3, 3)
).astype(np.float32)

W = np.array(
    [
        [
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],  # (1, 2, 3, 3)
            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        ]
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "ConvTranspose", ["X", "W"], ["Y"], strides=[3, 2], pads=[1, 2, 1, 2]
)

y = np.array(
    [
        [
            [
                [1.0, 1.0, 3.0],  # (1, 2, 7, 3)
                [1.0, 1.0, 3.0],
                [7.0, 4.0, 9.0],
                [7.0, 4.0, 9.0],
                [7.0, 4.0, 9.0],
                [13.0, 7.0, 15.0],
                [13.0, 7.0, 15.0],
            ],
            [
                [1.0, 1.0, 3.0],
                [1.0, 1.0, 3.0],
                [7.0, 4.0, 9.0],
                [7.0, 4.0, 9.0],
                [7.0, 4.0, 9.0],
                [13.0, 7.0, 15.0],
                [13.0, 7.0, 15.0],
            ],
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x, W], outputs=[y], name="test_convtranspose_pads")

Cos

There are 1 test cases, listed as following:

cos

node = onnx.helper.make_node(
    "Cos",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.cos(x)
expect(node, inputs=[x], outputs=[y], name="test_cos_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.cos(x)
expect(node, inputs=[x], outputs=[y], name="test_cos")

Cosh

There are 1 test cases, listed as following:

cosh

node = onnx.helper.make_node(
    "Cosh",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.cosh(x)  # expected output [1.54308069,  1.,  1.54308069]
expect(node, inputs=[x], outputs=[y], name="test_cosh_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.cosh(x)
expect(node, inputs=[x], outputs=[y], name="test_cosh")

CumProd

There are 9 test cases, listed as following:

cumprod_1d

node = onnx.helper.make_node("CumProd", inputs=["x", "axis"], outputs=["y"])
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.array(0, dtype=np.int32)
y = np.array([1.0, 2.0, 6.0, 24.0, 120.0]).astype(np.float64)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumprod_1d")

cumprod_1d_exclusive

node = onnx.helper.make_node(
    "CumProd", inputs=["x", "axis"], outputs=["y"], exclusive=1
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.array(0, dtype=np.int32)
y = np.array([1.0, 1.0, 2.0, 6.0, 24.0]).astype(np.float64)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumprod_1d_exclusive")

cumprod_1d_int32_exclusive

node = onnx.helper.make_node(
    "CumProd", inputs=["x", "axis"], outputs=["y"], exclusive=1
)
x = np.array([1, 2, 3, 4, 5]).astype(np.int32)
axis = np.array(0, dtype=np.int32)
y = np.array([1, 1, 2, 6, 24]).astype(np.int32)
expect(
    node, inputs=[x, axis], outputs=[y], name="test_cumprod_1d_int32_exclusive"
)

cumprod_1d_reverse

node = onnx.helper.make_node(
    "CumProd", inputs=["x", "axis"], outputs=["y"], reverse=1
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.array(0, dtype=np.int32)
y = np.array([120.0, 120.0, 60.0, 20.0, 5.0]).astype(np.float64)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumprod_1d_reverse")

cumprod_1d_reverse_exclusive

node = onnx.helper.make_node(
    "CumProd", inputs=["x", "axis"], outputs=["y"], reverse=1, exclusive=1
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.array(0, dtype=np.int32)
y = np.array([120.0, 60.0, 20.0, 5.0, 1.0]).astype(np.float64)
expect(
    node,
    inputs=[x, axis],
    outputs=[y],
    name="test_cumprod_1d_reverse_exclusive",
)

cumprod_2d_axis_0

node = onnx.helper.make_node(
    "CumProd",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float64).reshape((2, 3))
axis = np.array(0, dtype=np.int32)
y = (
    np.array([1.0, 2.0, 3.0, 4.0, 10.0, 18.0])
    .astype(np.float64)
    .reshape((2, 3))
)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumprod_2d_axis_0")

cumprod_2d_axis_1

node = onnx.helper.make_node(
    "CumProd",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float64).reshape((2, 3))
axis = np.array(1, dtype=np.int32)
y = (
    np.array([1.0, 2.0, 6.0, 4.0, 20.0, 120.0])
    .astype(np.float64)
    .reshape((2, 3))
)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumprod_2d_axis_1")

cumprod_2d_int32

node = onnx.helper.make_node(
    "CumProd",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1, 2, 3, 4, 5, 6]).astype(np.int32).reshape((2, 3))
axis = np.array(0, dtype=np.int32)
y = np.array([1, 2, 3, 4, 10, 18]).astype(np.int32).reshape((2, 3))
expect(node, inputs=[x, axis], outputs=[y], name="test_cumprod_2d_int32")

cumprod_2d_negative_axis

node = onnx.helper.make_node(
    "CumProd",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float64).reshape((2, 3))
axis = np.array(-1, dtype=np.int32)
y = (
    np.array([1.0, 2.0, 6.0, 4.0, 20.0, 120.0])
    .astype(np.float64)
    .reshape((2, 3))
)
expect(
    node, inputs=[x, axis], outputs=[y], name="test_cumprod_2d_negative_axis"
)

CumSum

There are 9 test cases, listed as following:

cumsum_1d

node = onnx.helper.make_node("CumSum", inputs=["x", "axis"], outputs=["y"])
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.int32(0)
y = np.array([1.0, 3.0, 6.0, 10.0, 15.0]).astype(np.float64)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_1d")

cumsum_1d_exclusive

node = onnx.helper.make_node(
    "CumSum", inputs=["x", "axis"], outputs=["y"], exclusive=1
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.int32(0)
y = np.array([0.0, 1.0, 3.0, 6.0, 10.0]).astype(np.float64)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_1d_exclusive")

cumsum_1d_int32_exclusive

node = onnx.helper.make_node(
    "CumSum", inputs=["x", "axis"], outputs=["y"], exclusive=1
)
x = np.array([1, 2, 3, 4, 5]).astype(np.int32)
axis = np.int32(0)
y = np.array([0, 1, 3, 6, 10]).astype(np.int32)
expect(
    node, inputs=[x, axis], outputs=[y], name="test_cumsum_1d_int32_exclusive"
)

cumsum_1d_reverse

node = onnx.helper.make_node(
    "CumSum", inputs=["x", "axis"], outputs=["y"], reverse=1
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.int32(0)
y = np.array([15.0, 14.0, 12.0, 9.0, 5.0]).astype(np.float64)
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_1d_reverse")

cumsum_1d_reverse_exclusive

node = onnx.helper.make_node(
    "CumSum", inputs=["x", "axis"], outputs=["y"], reverse=1, exclusive=1
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype(np.float64)
axis = np.int32(0)
y = np.array([14.0, 12.0, 9.0, 5.0, 0.0]).astype(np.float64)
expect(
    node, inputs=[x, axis], outputs=[y], name="test_cumsum_1d_reverse_exclusive"
)

cumsum_2d_axis_0

node = onnx.helper.make_node(
    "CumSum",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float64).reshape((2, 3))
axis = np.int32(0)
y = np.array([1.0, 2.0, 3.0, 5.0, 7.0, 9.0]).astype(np.float64).reshape((2, 3))
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_2d_axis_0")

cumsum_2d_axis_1

node = onnx.helper.make_node(
    "CumSum",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float64).reshape((2, 3))
axis = np.int32(1)
y = np.array([1.0, 3.0, 6.0, 4.0, 9.0, 15.0]).astype(np.float64).reshape((2, 3))
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_2d_axis_1")

cumsum_2d_int32

node = onnx.helper.make_node(
    "CumSum",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1, 2, 3, 4, 5, 6]).astype(np.int32).reshape((2, 3))
axis = np.int32(0)
y = np.array([1, 2, 3, 5, 7, 9]).astype(np.int32).reshape((2, 3))
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_2d_int32")

cumsum_2d_negative_axis

node = onnx.helper.make_node(
    "CumSum",
    inputs=["x", "axis"],
    outputs=["y"],
)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float64).reshape((2, 3))
axis = np.int32(-1)
y = np.array([1.0, 3.0, 6.0, 4.0, 9.0, 15.0]).astype(np.float64).reshape((2, 3))
expect(node, inputs=[x, axis], outputs=[y], name="test_cumsum_2d_negative_axis")

DFT

There are 2 test cases, listed as following:

dft

node = onnx.helper.make_node("DFT", inputs=["x", "", "axis"], outputs=["y"])
x = np.arange(0, 100).reshape(10, 10).astype(np.float32)
axis = np.array(1, dtype=np.int64)
y = np.fft.fft(x, axis=0)

x = x.reshape(1, 10, 10, 1)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
expect(node, inputs=[x, axis], outputs=[y], name="test_dft")

node = onnx.helper.make_node("DFT", inputs=["x", "", "axis"], outputs=["y"])
x = np.arange(0, 100).reshape(10, 10).astype(np.float32)
axis = np.array(2, dtype=np.int64)
y = np.fft.fft(x, axis=1)

x = x.reshape(1, 10, 10, 1)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
expect(node, inputs=[x, axis], outputs=[y], name="test_dft_axis")

node = onnx.helper.make_node(
    "DFT", inputs=["x", "", "axis"], outputs=["y"], inverse=1
)
x = np.arange(0, 100, dtype=np.complex64).reshape(10, 10)
axis = np.array(1, dtype=np.int64)
y = np.fft.ifft(x, axis=0)

x = np.stack((x.real, x.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
expect(node, inputs=[x, axis], outputs=[y], name="test_dft_inverse")

# Test RFFT (Real FFT): real input -> one-sided complex output
node = onnx.helper.make_node(
    "DFT", inputs=["x", "", "axis"], outputs=["y"], onesided=1
)
x = np.arange(0, 100).reshape(10, 10).astype(np.float32)
axis = np.array(1, dtype=np.int64)
y = np.fft.rfft(x, axis=0)

x = x.reshape(1, 10, 10, 1)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 6, 10, 2)
expect(node, inputs=[x, axis], outputs=[y], name="test_dft_rfft")

# Test IRFFT (Inverse Real FFT): one-sided complex input -> real output
node = onnx.helper.make_node(
    "DFT", inputs=["x", "", "axis"], outputs=["y"], onesided=1, inverse=1
)
# Create one-sided complex input (6 bins for signal length 10)
x = np.fft.rfft(np.arange(0, 100).reshape(10, 10), axis=0).astype(np.complex64)
axis = np.array(1, dtype=np.int64)
y = np.fft.irfft(x, n=10, axis=0)

x = np.stack((x.real, x.imag), axis=2).astype(np.float32).reshape(1, 6, 10, 2)
y = y.reshape(1, 10, 10, 1).astype(np.float32)
expect(node, inputs=[x, axis], outputs=[y], name="test_dft_irfft")

opset19

node = onnx.helper.make_node("DFT", inputs=["x"], outputs=["y"], axis=1)
x = np.arange(0, 100).reshape(10, 10).astype(np.float32)
y = np.fft.fft(x, axis=0)

x = x.reshape(1, 10, 10, 1)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dft_opset19",
    opset_imports=[onnx.helper.make_opsetid("", 19)],
)

node = onnx.helper.make_node("DFT", inputs=["x"], outputs=["y"], axis=2)
x = np.arange(0, 100).reshape(10, 10).astype(np.float32)
y = np.fft.fft(x, axis=1)

x = x.reshape(1, 10, 10, 1)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dft_axis_opset19",
    opset_imports=[onnx.helper.make_opsetid("", 19)],
)

node = onnx.helper.make_node(
    "DFT", inputs=["x"], outputs=["y"], inverse=1, axis=1
)
x = np.arange(0, 100, dtype=np.complex64).reshape(
    10,
    10,
)
y = np.fft.ifft(x, axis=0)

x = np.stack((x.real, x.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 10, 10, 2)
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dft_inverse_opset19",
    opset_imports=[onnx.helper.make_opsetid("", 19)],
)

# Test RFFT (Real FFT): real input -> one-sided complex output
node = onnx.helper.make_node(
    "DFT", inputs=["x"], outputs=["y"], onesided=1, axis=1
)
x = np.arange(0, 100).reshape(10, 10).astype(np.float32)
y = np.fft.rfft(x, axis=0)

x = x.reshape(1, 10, 10, 1)
y = np.stack((y.real, y.imag), axis=2).astype(np.float32).reshape(1, 6, 10, 2)
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dft_rfft_opset19",
    opset_imports=[onnx.helper.make_opsetid("", 19)],
)

# Test IRFFT (Inverse Real FFT): one-sided complex input -> real output
node = onnx.helper.make_node(
    "DFT", inputs=["x"], outputs=["y"], onesided=1, inverse=1, axis=1
)
# Create one-sided complex input (6 bins for signal length 10)
x = np.fft.rfft(np.arange(0, 100).reshape(10, 10), axis=0).astype(np.complex64)
y = np.fft.irfft(x, n=10, axis=0)

x = np.stack((x.real, x.imag), axis=2).astype(np.float32).reshape(1, 6, 10, 2)
y = y.reshape(1, 10, 10, 1).astype(np.float32)
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dft_irfft_opset19",
    opset_imports=[onnx.helper.make_opsetid("", 19)],
)

DeformConv

There are 3 test cases, listed as following:

deformconv

X = np.arange(9).astype(np.float32)
X.shape = (1, 1, 3, 3)
W = np.ones((1, 1, 2, 2), dtype=np.float32)

# Convolution with padding
offset_with_padding = np.zeros((1, 8, 4, 4), dtype=np.float32)
# h-coord of [0, 0] element of kernel, at output position [0, 0]
offset_with_padding[0, 0, 0, 0] = 0.5
# w-coord of [1, 0] element of kernel, at output position [1, 2]
offset_with_padding[0, 5, 1, 2] = -0.1

node_with_padding = onnx.helper.make_node(
    "DeformConv",
    inputs=["X", "W", "offset_with_padding"],
    outputs=["Y_with_padding"],
    kernel_shape=[2, 2],
    pads=[1, 1, 1, 1],
)
Y_with_padding = np.array(
    [
        [
            [
                [0.0, 1.0, 3.0, 2.0],  # (1, 1, 4, 4) output tensor
                [3.0, 8.0, 11.9, 7.0],
                [9.0, 20.0, 24.0, 13.0],
                [6.0, 13.0, 15.0, 8.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_with_padding,
    inputs=[X, W, offset_with_padding],
    outputs=[Y_with_padding],
    name="test_basic_deform_conv_with_padding",
)

# Convolution without padding
offset_without_padding = np.zeros((1, 8, 2, 2), dtype=np.float32)
# h-coord of [0, 0] element of kernel, at output position [0, 0]
offset_without_padding[0, 0, 0, 0] = 0.5
# w-coord of [1, 0] element of kernel, at output position [0, 1]
offset_without_padding[0, 5, 0, 1] = -0.1

node_without_padding = onnx.helper.make_node(
    "DeformConv",
    inputs=["X", "W", "offset_without_padding"],
    outputs=["Y_without_padding"],
    kernel_shape=[2, 2],
    pads=[0, 0, 0, 0],
)
Y_without_padding = np.array(
    [
        [
            [
                [9.5, 11.9],  # (1, 1, 2, 2) output tensor
                [20.0, 24.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node_without_padding,
    inputs=[X, W, offset_without_padding],
    outputs=[Y_without_padding],
    name="test_basic_deform_conv_without_padding",
)

deformconv_with_mask_bias

X = np.arange(9).astype(np.float32)
X.shape = (1, 1, 3, 3)
W = np.ones((1, 1, 2, 2), dtype=np.float32)
B = np.ones((1,), dtype=np.float32)

offset = np.zeros((1, 8, 2, 2), dtype=np.float32)
# h-coord of [0, 0] element of kernel, at output position [0, 0]
offset[0, 0, 0, 0] = 0.5
# w-coord of [1, 0] element of kernel, at output position [0, 1]
offset[0, 5, 0, 1] = -0.1

mask = np.ones((1, 4, 2, 2), dtype=np.float32)
mask[0, 2, 1, 1] = 0.2  # [1, 0] element of kernel at output position [1, 1]

node = onnx.helper.make_node(
    "DeformConv",
    inputs=["X", "W", "offset", "B", "mask"],
    outputs=["Y"],
    kernel_shape=[2, 2],
    pads=[0, 0, 0, 0],
)
Y = np.array(
    [
        [
            [
                [10.5, 12.9],  # (1, 1, 2, 2) output tensor
                [21.0, 19.4],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node,
    inputs=[X, W, offset, B, mask],
    outputs=[Y],
    name="test_deform_conv_with_mask_bias",
)

deformconv_with_multiple_offset_groups

X = np.zeros((1, 2, 3, 3), dtype=np.float32)
X[0, 0] = np.reshape(np.arange(9).astype(np.float32), (3, 3))
X[0, 1] = np.reshape(np.arange(8, -1, -1).astype(np.float32), (3, 3))
X.shape = (1, 2, 3, 3)
W = np.ones((1, 2, 2, 2), dtype=np.float32)

offset = np.zeros((1, 16, 2, 2), dtype=np.float32)
# h-coord of [0, 0] element of kernel in channel 0, at output position [0, 0]
offset[0, 0, 0, 0] = 0.5
# w-coord of [1, 0] element of kernel in channel 1, at output position [0, 1]
offset[0, 13, 0, 1] = -0.1

node = onnx.helper.make_node(
    "DeformConv",
    inputs=["X", "W", "offset"],
    outputs=["Y"],
    kernel_shape=[2, 2],
    pads=[0, 0, 0, 0],
    offset_group=2,
)
Y = np.array(
    [
        [
            [
                [33.5, 32.1],  # (1, 1, 2, 2) output tensor
                [32.0, 32.0],
            ]
        ]
    ]
).astype(np.float32)
expect(
    node,
    inputs=[X, W, offset],
    outputs=[Y],
    name="test_deform_conv_with_multiple_offset_groups",
)

DepthToSpace

There are 2 test cases, listed as following:

crd_mode_example

node = onnx.helper.make_node(
    "DepthToSpace", inputs=["x"], outputs=["y"], blocksize=2, mode="CRD"
)

# (1, 8, 2, 3) input tensor
x = np.array(
    [
        [
            [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]],
            [[9.0, 10.0, 11.0], [12.0, 13.0, 14.0]],
            [[18.0, 19.0, 20.0], [21.0, 22.0, 23.0]],
            [[27.0, 28.0, 29.0], [30.0, 31.0, 32.0]],
            [[36.0, 37.0, 38.0], [39.0, 40.0, 41.0]],
            [[45.0, 46.0, 47.0], [48.0, 49.0, 50.0]],
            [[54.0, 55.0, 56.0], [57.0, 58.0, 59.0]],
            [[63.0, 64.0, 65.0], [66.0, 67.0, 68.0]],
        ]
    ]
).astype(np.float32)

# (1, 2, 4, 6) output tensor
y = np.array(
    [
        [
            [
                [0.0, 9.0, 1.0, 10.0, 2.0, 11.0],
                [18.0, 27.0, 19.0, 28.0, 20.0, 29.0],
                [3.0, 12.0, 4.0, 13.0, 5.0, 14.0],
                [21.0, 30.0, 22.0, 31.0, 23.0, 32.0],
            ],
            [
                [36.0, 45.0, 37.0, 46.0, 38.0, 47.0],
                [54.0, 63.0, 55.0, 64.0, 56.0, 65.0],
                [39.0, 48.0, 40.0, 49.0, 41.0, 50.0],
                [57.0, 66.0, 58.0, 67.0, 59.0, 68.0],
            ],
        ]
    ]
).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_depthtospace_crd_mode_example")

default_mode_example

node = onnx.helper.make_node(
    "DepthToSpace", inputs=["x"], outputs=["y"], blocksize=2, mode="DCR"
)

# (1, 8, 2, 3) input tensor
x = np.array(
    [
        [
            [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]],
            [[9.0, 10.0, 11.0], [12.0, 13.0, 14.0]],
            [[18.0, 19.0, 20.0], [21.0, 22.0, 23.0]],
            [[27.0, 28.0, 29.0], [30.0, 31.0, 32.0]],
            [[36.0, 37.0, 38.0], [39.0, 40.0, 41.0]],
            [[45.0, 46.0, 47.0], [48.0, 49.0, 50.0]],
            [[54.0, 55.0, 56.0], [57.0, 58.0, 59.0]],
            [[63.0, 64.0, 65.0], [66.0, 67.0, 68.0]],
        ]
    ]
).astype(np.float32)

# (1, 2, 4, 6) output tensor
y = np.array(
    [
        [
            [
                [0.0, 18.0, 1.0, 19.0, 2.0, 20.0],
                [36.0, 54.0, 37.0, 55.0, 38.0, 56.0],
                [3.0, 21.0, 4.0, 22.0, 5.0, 23.0],
                [39.0, 57.0, 40.0, 58.0, 41.0, 59.0],
            ],
            [
                [9.0, 27.0, 10.0, 28.0, 11.0, 29.0],
                [45.0, 63.0, 46.0, 64.0, 47.0, 65.0],
                [12.0, 30.0, 13.0, 31.0, 14.0, 32.0],
                [48.0, 66.0, 49.0, 67.0, 50.0, 68.0],
            ],
        ]
    ]
).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_depthtospace_example")

DequantizeLinear

There are 14 test cases, listed as following:

axis

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
)

# 1-D tensor zero point and scale of size equal to axis 1 of the input tensor
x = np.array(
    [
        [
            [[3, 89], [34, 200], [74, 59]],
            [[5, 24], [24, 87], [32, 13]],
            [[245, 99], [4, 142], [121, 102]],
        ],
    ],
    dtype=np.uint8,
)
x_scale = np.array([2, 4, 5], dtype=np.float32)
x_zero_point = np.array([84, 24, 196], dtype=np.uint8)
y = (
    x.astype(np.float32) - x_zero_point.reshape(1, 3, 1, 1).astype(np.float32)
) * x_scale.reshape(1, 3, 1, 1)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_axis",
)

blocked

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
    axis=1,
    block_size=2,
)

x = np.array(
    [
        [
            [[3, 89], [34, 200], [74, 59]],
            [[5, 24], [24, 87], [32, 13]],
            [[5, 12], [12, 33], [65, 42]],
            [[245, 99], [4, 142], [121, 102]],
        ],
    ],
    dtype=np.uint8,
)

x_scale = np.array(
    [
        [
            [[3.0, 2.0], [4.0, 1.0], [2.0, 2.0]],
            [[5.0, 2.0], [4.0, 3.0], [5.0, 2.0]],
        ],
    ],
    dtype=np.float32,
)
x_zero_point = np.array(
    [
        [
            [[1, 0], [0, 1], [2, 20]],
            [[3, 2], [4, 3], [15, 2]],
        ],
    ],
    dtype=np.uint8,
)

# x.shape = (1, 4, 3, 2)
# x_scale.shape = (1, 2, 3, 2)
assert x_scale.shape == x_zero_point.shape
block_axis = 1
# The block shape is [x.shape[i] // x_scale.shape[i] for i in range(len(x.shape))] = (1, 2, 1, 1)
assert all(
    x.shape[i] == x_scale.shape[i]
    for i in range(len(x.shape))
    if i != block_axis
)
assert x.shape[block_axis] % x_scale.shape[block_axis] == 0
repeats = x.shape[block_axis] // x_scale.shape[block_axis]

# Create element-wise scale and zero point
x_scale_elementwise = np.repeat(x_scale, repeats=repeats, axis=block_axis)
x_zero_point_elementwise = np.repeat(
    x_zero_point, repeats=repeats, axis=block_axis
)

y = (
    x.astype(np.float32) - x_zero_point_elementwise.astype(np.float32)
) * x_scale_elementwise

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_blocked",
)

dequantizelinear

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
)

# scalar zero point and scale
x = np.array([0, 3, 128, 255]).astype(np.uint8)
x_scale = np.float32(2)
x_zero_point = np.uint8(128)
y = np.array([-256, -250, 0, 254], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear",
)

e4m3fn

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.FLOAT8E4M3FN, [5], [0, 0.5, 1, 448, -104])
x_scale = np.float32(2)
y = np.array([0.0, 1.0, 2.0, 896.0, -208.0], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale],
    outputs=[y],
    name="test_dequantizelinear_e4m3fn",
)

e4m3fn_float16

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.FLOAT8E4M3FN, [5], [0, 0.5, 1, 448, -104])
x_scale = np.float16(2)
y = np.array([0.0, 1.0, 2.0, 896.0, -208.0], dtype=np.float16)

expect(
    node,
    inputs=[x, x_scale],
    outputs=[y],
    name="test_dequantizelinear_e4m3fn_float16",
)

e4m3fn_zero_point

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "zero_point"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.FLOAT8E4M3FN, [5], [0, 0.5, 1, 448, -104])
zero_point = make_tensor("zero_point", TensorProto.FLOAT8E4M3FN, [1], [0])
x_scale = np.float32(2)
y = np.array([0.0, 1.0, 2.0, 896.0, -208.0], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, zero_point],
    outputs=[y],
    name="test_dequantizelinear_e4m3fn_zero_point",
)

e5m2

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.FLOAT8E5M2, [5], [0, 0.5, 1, 49152, -96])
x_scale = np.float32(2)
y = np.array([0.0, 1.0, 2.0, 98304.0, -192.0], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale],
    outputs=[y],
    name="test_dequantizelinear_e5m2",
)

float4e2m1

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.FLOAT4E2M1, [5], [0, 1, -1, 1.5, -4])
x_scale = np.float32(2)
x_zero_point = make_tensor("x_zero_point", TensorProto.FLOAT4E2M1, (1,), [0])
y = np.array([0, 2, -2, 3, -8], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_float4e2m1",
)

int16

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
)

x = np.array([-300, -30, -1025, 1270]).astype(np.int16)
x_scale = np.float32(2)
x_zero_point = np.int16(-1024)
y = np.array([1448.0, 1988.0, -2.0, 4588.0], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_int16",
)

int2

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.INT2, [4], [0, 1, -1, -2])
x_scale = np.float32(2)
x_zero_point = make_tensor("x_zero_point", TensorProto.INT2, (1,), [1])
y = np.array([-2, 0, -4, -6], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_int2",
)

int4

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.INT4, [5], [0, 1, 7, -4, -8])
x_scale = np.float32(2)
x_zero_point = make_tensor("x_zero_point", TensorProto.INT4, (1,), [1])
y = np.array([-2, 0, 12, -10, -18], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_int4",
)

uint16

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
)

x = np.array([30000, 31000, 32768, 33000]).astype(np.uint16)
x_scale = np.float32(2)
x_zero_point = np.uint16(32767)
y = np.array([-5534.0, -3534.0, 2.0, 466.0], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_uint16",
)

uint2

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.UINT2, [4], [0, 1, 2, 3])
x_scale = np.float32(2)
x_zero_point = make_tensor("x_zero_point", TensorProto.UINT2, (1,), [1])
y = np.array([-2, 0, 2, 4], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_uint2",
)

uint4

node = onnx.helper.make_node(
    "DequantizeLinear",
    inputs=["x", "x_scale", "x_zero_point"],
    outputs=["y"],
    axis=0,
)

# scalar zero point and scale
x = make_tensor("x", TensorProto.UINT4, [5], [0, 1, 7, 10, 15])
x_scale = np.float32(2)
x_zero_point = make_tensor("x_zero_point", TensorProto.UINT4, (1,), [1])
y = np.array([-2, 0, 12, 18, 28], dtype=np.float32)

expect(
    node,
    inputs=[x, x_scale, x_zero_point],
    outputs=[y],
    name="test_dequantizelinear_uint4",
)

Det

There are 2 test cases, listed as following:

2d

node = onnx.helper.make_node(
    "Det",
    inputs=["x"],
    outputs=["y"],
)

x = np.arange(4).reshape(2, 2).astype(np.float32)
y = np.linalg.det(x)  # expect -2
expect(node, inputs=[x], outputs=[y], name="test_det_2d")

nd

node = onnx.helper.make_node(
    "Det",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]).astype(
    np.float32
)
y = np.linalg.det(x)  # expect array([-2., -3., -8.])
expect(node, inputs=[x], outputs=[y], name="test_det_nd")

Div

There are 2 test cases, listed as following:

div

node = onnx.helper.make_node(
    "Div",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([3, 4]).astype(np.float32)
y = np.array([1, 2]).astype(np.float32)
z = x / y  # expected output [3., 2.]
expect(node, inputs=[x, y], outputs=[z], name="test_div_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.rand(3, 4, 5).astype(np.float32) + 1.0
z = x / y
expect(node, inputs=[x, y], outputs=[z], name="test_div")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.int8) + 1
z = x // y
expect(node, inputs=[x, y], outputs=[z], name="test_div_int8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.int16) + 1
z = x // y
expect(node, inputs=[x, y], outputs=[z], name="test_div_int16")

x = np.array([-3, 3, -3, 3], dtype=np.int32)
y = np.array([2, 2, -2, -2], dtype=np.int32)
z = np.array([-1, 1, 1, -1], dtype=np.int32)
expect(node, inputs=[x, y], outputs=[z], name="test_div_int32_trunc")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8) + 1
z = x // y
expect(node, inputs=[x, y], outputs=[z], name="test_div_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16) + 1
z = x // y
expect(node, inputs=[x, y], outputs=[z], name="test_div_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32) + 1
z = x // y
expect(node, inputs=[x, y], outputs=[z], name="test_div_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64) + 1
z = x // y
expect(node, inputs=[x, y], outputs=[z], name="test_div_uint64")

div_broadcast

node = onnx.helper.make_node(
    "Div",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.rand(5).astype(np.float32) + 1.0
z = x / y
expect(node, inputs=[x, y], outputs=[z], name="test_div_bcast")

Dropout

There are 12 test cases, listed as following:

default

seed = np.int64(0)
node = onnx.helper.make_node("Dropout", inputs=["x"], outputs=["y"], seed=seed)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = dropout(x)
expect(node, inputs=[x], outputs=[y], name="test_dropout_default")

default_mask

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x"], outputs=["y", "z"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y, z = dropout(x, return_mask=True)
expect(node, inputs=[x], outputs=[y, z], name="test_dropout_default_mask")

default_mask_ratio

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r"], outputs=["y", "z"], seed=seed
)

r = np.float32(0.1)
x = np.random.randn(3, 4, 5).astype(np.float32)
y, z = dropout(x, r, return_mask=True)
expect(
    node, inputs=[x, r], outputs=[y, z], name="test_dropout_default_mask_ratio"
)

default_old

node = onnx.helper.make_node(
    "Dropout",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = x
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dropout_default_old",
    opset_imports=[helper.make_opsetid("", 11)],
)

default_ratio

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r"], outputs=["y"], seed=seed
)

r = np.float32(0.1)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = dropout(x, r)
expect(node, inputs=[x, r], outputs=[y], name="test_dropout_default_ratio")

random_old

node = onnx.helper.make_node(
    "Dropout",
    inputs=["x"],
    outputs=["y"],
    ratio=0.2,
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = x
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_dropout_random_old",
    opset_imports=[helper.make_opsetid("", 11)],
)

training

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r", "t"], outputs=["y"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
r = np.float32(0.75)
t = np.bool_(True)
y = dropout(x, r, training_mode=t)
expect(node, inputs=[x, r, t], outputs=[y], name="test_training_dropout")

training_default

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r", "t"], outputs=["y"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
r = np.float32(0.5)
t = np.bool_(True)
y = dropout(x, r, training_mode=t)
expect(
    node, inputs=[x, r, t], outputs=[y], name="test_training_dropout_default"
)

training_default_ratio_mask

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r", "t"], outputs=["y", "z"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
r = np.float32(0.5)
t = np.bool_(True)
y, z = dropout(x, r, training_mode=t, return_mask=True)
expect(
    node,
    inputs=[x, r, t],
    outputs=[y, z],
    name="test_training_dropout_default_mask",
)

training_default_zero_ratio

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r", "t"], outputs=["y"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
r = np.float32(0.0)
t = np.bool_(True)
y = dropout(x, r, training_mode=t)
expect(
    node, inputs=[x, r, t], outputs=[y], name="test_training_dropout_zero_ratio"
)

training_default_zero_ratio_mask

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r", "t"], outputs=["y", "z"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
r = np.float32(0.0)
t = np.bool_(True)
y, z = dropout(x, r, training_mode=t, return_mask=True)
expect(
    node,
    inputs=[x, r, t],
    outputs=[y, z],
    name="test_training_dropout_zero_ratio_mask",
)

training_ratio_mask

seed = np.int64(0)
node = onnx.helper.make_node(
    "Dropout", inputs=["x", "r", "t"], outputs=["y", "z"], seed=seed
)

x = np.random.randn(3, 4, 5).astype(np.float32)
r = np.float32(0.75)
t = np.bool_(True)
y, z = dropout(x, r, training_mode=t, return_mask=True)
expect(
    node, inputs=[x, r, t], outputs=[y, z], name="test_training_dropout_mask"
)

DynamicQuantizeLinear

There are 1 test cases, listed as following:

dynamicquantizelinear

node = onnx.helper.make_node(
    "DynamicQuantizeLinear",
    inputs=["x"],
    outputs=["y", "y_scale", "y_zero_point"],
)

# expected scale 0.0196078438 and zero point 153
X = np.array([0, 2, -3, -2.5, 1.34, 0.5]).astype(np.float32)
x_min = np.minimum(0, np.min(X))
x_max = np.maximum(0, np.max(X))
Y_Scale = np.float32((x_max - x_min) / (255 - 0))  # uint8 -> [0, 255]
Y_ZeroPoint = np.clip(round((0 - x_min) / Y_Scale), 0, 255).astype(np.uint8)
Y = np.clip(np.round(X / Y_Scale) + Y_ZeroPoint, 0, 255).astype(np.uint8)

expect(
    node,
    inputs=[X],
    outputs=[Y, Y_Scale, Y_ZeroPoint],
    name="test_dynamicquantizelinear",
)

# expected scale 0.0156862754 and zero point 255
X = np.array([-1.0, -2.1, -1.3, -2.5, -3.34, -4.0]).astype(np.float32)
x_min = np.minimum(0, np.min(X))
x_max = np.maximum(0, np.max(X))
Y_Scale = np.float32((x_max - x_min) / (255 - 0))  # uint8 -> [0, 255]
Y_ZeroPoint = np.clip(round((0 - x_min) / Y_Scale), 0, 255).astype(np.uint8)
Y = np.clip(np.round(X / Y_Scale) + Y_ZeroPoint, 0, 255).astype(np.uint8)

expect(
    node,
    inputs=[X],
    outputs=[Y, Y_Scale, Y_ZeroPoint],
    name="test_dynamicquantizelinear_max_adjusted",
)

X = (
    np.array([1, 2.1, 1.3, 2.5, 3.34, 4.0, 1.5, 2.6, 3.9, 4.0, 3.0, 2.345])
    .astype(np.float32)
    .reshape((3, 4))
)

# expected scale 0.0156862754 and zero point 0
x_min = np.minimum(0, np.min(X))
x_max = np.maximum(0, np.max(X))
Y_Scale = np.float32((x_max - x_min) / (255 - 0))  # uint8 -> [0, 255]
Y_ZeroPoint = np.clip(round((0 - x_min) / Y_Scale), 0, 255).astype(np.uint8)
Y = np.clip(np.round(X / Y_Scale) + Y_ZeroPoint, 0, 255).astype(np.uint8)

expect(
    node,
    inputs=[X],
    outputs=[Y, Y_Scale, Y_ZeroPoint],
    name="test_dynamicquantizelinear_min_adjusted",
)

Einsum

There are 6 test cases, listed as following:

einsum_batch_diagonal

Eqn = "...ii ->...i"
node = onnx.helper.make_node(
    "Einsum", inputs=["x"], outputs=["y"], equation=Eqn
)

X = np.random.randn(3, 5, 5)
Z = einsum_reference_implementation(Eqn, (X,))

expect(node, inputs=[X], outputs=[Z], name="test_einsum_batch_diagonal")

einsum_batch_matmul

Eqn = "bij, bjk -> bik"
node = onnx.helper.make_node(
    "Einsum", inputs=["x", "y"], outputs=["z"], equation=Eqn
)

X = np.random.randn(5, 2, 3)
Y = np.random.randn(5, 3, 4)
Z = einsum_reference_implementation(Eqn, (X, Y))

expect(node, inputs=[X, Y], outputs=[Z], name="test_einsum_batch_matmul")

einsum_inner_prod

Eqn = "i,i"
node = onnx.helper.make_node(
    "Einsum", inputs=["x", "y"], outputs=["z"], equation=Eqn
)

X = np.random.randn(5)
Y = np.random.randn(5)
Z = einsum_reference_implementation(Eqn, (X, Y))

expect(node, inputs=[X, Y], outputs=[Z], name="test_einsum_inner_prod")

einsum_scalar

Eqn = "->"
node = onnx.helper.make_node(
    "Einsum", inputs=["x"], outputs=["y"], equation=Eqn
)

X = np.array(5.0)  # scalar input
Z = einsum_reference_implementation(Eqn, (X,))

expect(node, inputs=[X], outputs=[Z], name="test_einsum_scalar")

einsum_sum

Eqn = "ij->i"
node = onnx.helper.make_node(
    "Einsum", inputs=["x"], outputs=["y"], equation=Eqn
)

X = np.random.randn(3, 4)
Z = einsum_reference_implementation(Eqn, (X,))

expect(node, inputs=[X], outputs=[Z], name="test_einsum_sum")

einsum_transpose

Eqn = "ij->ji"
node = onnx.helper.make_node(
    "Einsum", inputs=["x"], outputs=["y"], equation=Eqn
)

X = np.random.randn(3, 4)
Y = einsum_reference_implementation(Eqn, (X,))

expect(node, inputs=[X], outputs=[Y], name="test_einsum_transpose")

Elu

There are 2 test cases, listed as following:

elu

node = onnx.helper.make_node("Elu", inputs=["x"], outputs=["y"], alpha=2.0)

x = np.array([-1, 0, 1]).astype(np.float32)
# expected output [-1.2642411, 0., 1.]
y = np.clip(x, 0, np.inf) + (np.exp(np.clip(x, -np.inf, 0)) - 1) * 2.0
expect(node, inputs=[x], outputs=[y], name="test_elu_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, 0, np.inf) + (np.exp(np.clip(x, -np.inf, 0)) - 1) * 2.0
expect(node, inputs=[x], outputs=[y], name="test_elu")

elu_default

default_alpha = 1.0
node = onnx.helper.make_node(
    "Elu",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, 0, np.inf) + (np.exp(np.clip(x, -np.inf, 0)) - 1) * default_alpha
expect(node, inputs=[x], outputs=[y], name="test_elu_default")

Equal

There are 4 test cases, listed as following:

equal

node = onnx.helper.make_node(
    "Equal",
    inputs=["x", "y"],
    outputs=["z"],
)

x = (np.random.randn(3, 4, 5) * 10).astype(np.int32)
y = (np.random.randn(3, 4, 5) * 10).astype(np.int32)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal")

x = (np.random.randn(3, 4, 5) * 10).astype(np.int8)
y = (np.random.randn(3, 4, 5) * 10).astype(np.int8)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_int8")

x = (np.random.randn(3, 4, 5) * 10).astype(np.int16)
y = (np.random.randn(3, 4, 5) * 10).astype(np.int16)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_int16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_uint64")

equal_broadcast

node = onnx.helper.make_node(
    "Equal",
    inputs=["x", "y"],
    outputs=["z"],
)

x = (np.random.randn(3, 4, 5) * 10).astype(np.int32)
y = (np.random.randn(5) * 10).astype(np.int32)
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_bcast")

equal_string

node = onnx.helper.make_node(
    "Equal",
    inputs=["x", "y"],
    outputs=["z"],
)
x = np.array(["string1", "string2"], dtype=np.dtype(object))
y = np.array(["string1", "string3"], dtype=np.dtype(object))
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_string")

equal_string_broadcast

node = onnx.helper.make_node(
    "Equal",
    inputs=["x", "y"],
    outputs=["z"],
)
x = np.array(["string1", "string2"], dtype=np.dtype(object))
y = np.array(["string1"], dtype=np.dtype(object))
z = np.equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_equal_string_broadcast")

Erf

There are 1 test cases, listed as following:

erf

node = onnx.helper.make_node(
    "Erf",
    inputs=["x"],
    outputs=["y"],
)

x = np.random.randn(1, 3, 32, 32).astype(np.float32)
y = np.vectorize(math.erf)(x).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_erf")

Exp

There are 1 test cases, listed as following:

exp

node = onnx.helper.make_node(
    "Exp",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.exp(x)  # expected output [0.36787945, 1., 2.71828175]
expect(node, inputs=[x], outputs=[y], name="test_exp_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.exp(x)
expect(node, inputs=[x], outputs=[y], name="test_exp")

Expand

There are 2 test cases, listed as following:

dim_changed

node = onnx.helper.make_node(
    "Expand",
    inputs=["data", "new_shape"],
    outputs=["expanded"],
)
shape = [3, 1]
data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[1.], [2.], [3.]]
new_shape = [2, 1, 6]
expanded = data * np.ones(new_shape, dtype=np.float32)
# print(expanded)
# [[[1., 1., 1., 1., 1., 1.],
#  [2., 2., 2., 2., 2., 2.],
#  [3., 3., 3., 3., 3., 3.]],
#
# [[1., 1., 1., 1., 1., 1.],
#  [2., 2., 2., 2., 2., 2.],
#  [3., 3., 3., 3., 3., 3.]]]
new_shape = np.array(new_shape, dtype=np.int64)
expect(
    node,
    inputs=[data, new_shape],
    outputs=[expanded],
    name="test_expand_dim_changed",
)

dim_unchanged

node = onnx.helper.make_node(
    "Expand",
    inputs=["data", "new_shape"],
    outputs=["expanded"],
)
shape = [3, 1]
new_shape = [3, 4]
data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[1.], [2.], [3.]]
expanded = np.tile(data, 4)
# print(expanded)
# [[1., 1., 1., 1.],
# [2., 2., 2., 2.],
# [3., 3., 3., 3.]]
new_shape = np.array(new_shape, dtype=np.int64)
expect(
    node,
    inputs=[data, new_shape],
    outputs=[expanded],
    name="test_expand_dim_unchanged",
)

EyeLike

There are 3 test cases, listed as following:

populate_off_main_diagonal

shape = (4, 5)
off_diagonal_offset = 1
node = onnx.helper.make_node(
    "EyeLike",
    inputs=["x"],
    outputs=["y"],
    k=off_diagonal_offset,
    dtype=onnx.TensorProto.FLOAT,
)

x = np.random.randint(0, 100, size=shape, dtype=np.int32)
y = np.eye(shape[0], shape[1], k=off_diagonal_offset, dtype=np.float32)
expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_eyelike_populate_off_main_diagonal",
)

with_dtype

shape = (3, 4)
node = onnx.helper.make_node(
    "EyeLike",
    inputs=["x"],
    outputs=["y"],
    dtype=onnx.TensorProto.DOUBLE,
)

x = np.random.randint(0, 100, size=shape, dtype=np.int32)
y = np.eye(shape[0], shape[1], dtype=np.float64)
expect(node, inputs=[x], outputs=[y], name="test_eyelike_with_dtype")

without_dtype

shape = (4, 4)
node = onnx.helper.make_node(
    "EyeLike",
    inputs=["x"],
    outputs=["y"],
)

x = np.random.randint(0, 100, size=shape, dtype=np.int32)
y = np.eye(shape[0], shape[1], dtype=np.int32)
expect(node, inputs=[x], outputs=[y], name="test_eyelike_without_dtype")

Flatten

There are 3 test cases, listed as following:

flatten

shape = (2, 3, 4, 5)
a = np.random.random_sample(shape).astype(np.float32)

for i in range(len(shape)):
    node = onnx.helper.make_node(
        "Flatten",
        inputs=["a"],
        outputs=["b"],
        axis=i,
    )

    new_shape = (1, -1) if i == 0 else (np.prod(shape[0:i]).astype(int), -1)
    b = np.reshape(a, new_shape)
    expect(node, inputs=[a], outputs=[b], name="test_flatten_axis" + str(i))

flatten_negative_axis

shape = (2, 3, 4, 5)
a = np.random.random_sample(shape).astype(np.float32)

for i in range(-len(shape), 0):
    node = onnx.helper.make_node(
        "Flatten",
        inputs=["a"],
        outputs=["b"],
        axis=i,
    )

    new_shape = (np.prod(shape[0:i]).astype(int), -1)
    b = np.reshape(a, new_shape)
    expect(
        node,
        inputs=[a],
        outputs=[b],
        name="test_flatten_negative_axis" + str(abs(i)),
    )

flatten_with_default_axis

node = onnx.helper.make_node(
    "Flatten",
    inputs=["a"],
    outputs=["b"],  # Default value for axis: axis=1
)

shape = (5, 4, 3, 2)
a = np.random.random_sample(shape).astype(np.float32)
new_shape = (5, 24)
b = np.reshape(a, new_shape)
expect(node, inputs=[a], outputs=[b], name="test_flatten_default_axis")

Floor

There are 1 test cases, listed as following:

floor

node = onnx.helper.make_node(
    "Floor",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1.5, 1.2, 2]).astype(np.float32)
y = np.floor(x)  # expected output [-2., 1., 2.]
expect(node, inputs=[x], outputs=[y], name="test_floor_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.floor(x)
expect(node, inputs=[x], outputs=[y], name="test_floor")

GRU

There are 4 test cases, listed as following:

batchwise

input = np.array([[[1.0, 2.0]], [[3.0, 4.0]], [[5.0, 6.0]]]).astype(np.float32)

input_size = 2
hidden_size = 6
number_of_gates = 3
weight_scale = 0.2
layout = 1

node = onnx.helper.make_node(
    "GRU",
    inputs=["X", "W", "R"],
    outputs=["Y", "Y_h"],
    hidden_size=hidden_size,
    layout=layout,
)

W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)

gru = GRUHelper(X=input, W=W, R=R, layout=layout)
Y, Y_h = gru.step()
expect(
    node,
    inputs=[input, W, R],
    outputs=[Y.astype(np.float32), Y_h.astype(np.float32)],
    name="test_gru_batchwise",
)

defaults

input = np.array([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]).astype(np.float32)

input_size = 2
hidden_size = 5
weight_scale = 0.1
number_of_gates = 3

node = onnx.helper.make_node(
    "GRU", inputs=["X", "W", "R"], outputs=["", "Y_h"], hidden_size=hidden_size
)

W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)

gru = GRUHelper(X=input, W=W, R=R)
_, Y_h = gru.step()
expect(
    node,
    inputs=[input, W, R],
    outputs=[Y_h.astype(np.float32)],
    name="test_gru_defaults",
)

initial_bias

input = np.array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]).astype(
    np.float32
)

input_size = 3
hidden_size = 3
weight_scale = 0.1
custom_bias = 0.1
number_of_gates = 3

node = onnx.helper.make_node(
    "GRU",
    inputs=["X", "W", "R", "B"],
    outputs=["", "Y_h"],
    hidden_size=hidden_size,
)

W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)

# Adding custom bias
W_B = custom_bias * np.ones((1, number_of_gates * hidden_size)).astype(
    np.float32
)
R_B = np.zeros((1, number_of_gates * hidden_size)).astype(np.float32)
B = np.concatenate((W_B, R_B), axis=1)

gru = GRUHelper(X=input, W=W, R=R, B=B)
_, Y_h = gru.step()
expect(
    node,
    inputs=[input, W, R, B],
    outputs=[Y_h.astype(np.float32)],
    name="test_gru_with_initial_bias",
)

seq_length

input = np.array(
    [
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
        [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]],
    ]
).astype(np.float32)

input_size = 3
hidden_size = 5
number_of_gates = 3

node = onnx.helper.make_node(
    "GRU",
    inputs=["X", "W", "R", "B"],
    outputs=["", "Y_h"],
    hidden_size=hidden_size,
)

W = np.random.randn(1, number_of_gates * hidden_size, input_size).astype(
    np.float32
)
R = np.random.randn(1, number_of_gates * hidden_size, hidden_size).astype(
    np.float32
)

# Adding custom bias
W_B = np.random.randn(1, number_of_gates * hidden_size).astype(np.float32)
R_B = np.random.randn(1, number_of_gates * hidden_size).astype(np.float32)
B = np.concatenate((W_B, R_B), axis=1)

gru = GRUHelper(X=input, W=W, R=R, B=B)
_, Y_h = gru.step()
expect(
    node,
    inputs=[input, W, R, B],
    outputs=[Y_h.astype(np.float32)],
    name="test_gru_seq_length",
)

Gather

There are 4 test cases, listed as following:

gather_0

node = onnx.helper.make_node(
    "Gather",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=0,
)
data = np.random.randn(5, 4, 3, 2).astype(np.float32)
indices = np.array([0, 1, 3])
y = np.take(data, indices, axis=0)

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_0",
)

gather_1

node = onnx.helper.make_node(
    "Gather",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=1,
)
data = np.random.randn(5, 4, 3, 2).astype(np.float32)
indices = np.array([0, 1, 3])
y = np.take(data, indices, axis=1)

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_1",
)

gather_2d_indices

node = onnx.helper.make_node(
    "Gather",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=1,
)
data = np.random.randn(3, 3).astype(np.float32)
indices = np.array([[0, 2]])
y = np.take(data, indices, axis=1)

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_2d_indices",
)

gather_negative_indices

node = onnx.helper.make_node(
    "Gather",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=0,
)
data = np.arange(10).astype(np.float32)
indices = np.array([0, -9, -10])
y = np.take(data, indices, axis=0)

# print(y)
# [0. 1. 0.]

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_negative_indices",
)

GatherElements

There are 3 test cases, listed as following:

gather_elements_0

axis = 1
node = onnx.helper.make_node(
    "GatherElements",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=axis,
)
data = np.array([[1, 2], [3, 4]], dtype=np.float32)
indices = np.array([[0, 0], [1, 0]], dtype=np.int32)

y = gather_elements(data, indices, axis)
# print(y) produces
# [[1, 1],
#  [4, 3]]

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_elements_0",
)

gather_elements_1

axis = 0
node = onnx.helper.make_node(
    "GatherElements",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=axis,
)
data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
indices = np.array([[1, 2, 0], [2, 0, 0]], dtype=np.int32)

y = gather_elements(data, indices, axis)
# print(y) produces
# [[4, 8, 3],
#  [7, 2, 3]]

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_elements_1",
)

gather_elements_negative_indices

axis = 0
node = onnx.helper.make_node(
    "GatherElements",
    inputs=["data", "indices"],
    outputs=["y"],
    axis=axis,
)
data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
indices = np.array([[-1, -2, 0], [-2, 0, 0]], dtype=np.int32)

y = gather_elements(data, indices, axis)
# print(y) produces
# [[7, 5, 3],
#  [4, 2, 3]]

expect(
    node,
    inputs=[data, indices.astype(np.int64)],
    outputs=[y],
    name="test_gather_elements_negative_indices",
)

GatherND

There are 3 test cases, listed as following:

float32

node = onnx.helper.make_node(
    "GatherND",
    inputs=["data", "indices"],
    outputs=["output"],
)

data = np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=np.float32)
indices = np.array([[[0, 1]], [[1, 0]]], dtype=np.int64)
output = gather_nd_impl(data, indices, 0)
expected_output = np.array([[[2, 3]], [[4, 5]]], dtype=np.float32)
assert np.array_equal(output, expected_output)
expect(
    node,
    inputs=[data, indices],
    outputs=[output],
    name="test_gathernd_example_float32",
)

int32

node = onnx.helper.make_node(
    "GatherND",
    inputs=["data", "indices"],
    outputs=["output"],
)

data = np.array([[0, 1], [2, 3]], dtype=np.int32)
indices = np.array([[0, 0], [1, 1]], dtype=np.int64)
output = gather_nd_impl(data, indices, 0)
expected_output = np.array([0, 3], dtype=np.int32)
assert np.array_equal(output, expected_output)
expect(
    node,
    inputs=[data, indices],
    outputs=[output],
    name="test_gathernd_example_int32",
)

int32_batchdim_1

node = onnx.helper.make_node(
    "GatherND",
    inputs=["data", "indices"],
    outputs=["output"],
    batch_dims=1,
)

data = np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=np.int32)
indices = np.array([[1], [0]], dtype=np.int64)
output = gather_nd_impl(data, indices, 1)
expected_output = np.array([[2, 3], [4, 5]], dtype=np.int32)
assert np.array_equal(output, expected_output)
expect(
    node,
    inputs=[data, indices],
    outputs=[output],
    name="test_gathernd_example_int32_batch_dim1",
)

Gelu

There are 2 test cases, listed as following:

gelu_default

node = onnx.helper.make_node("Gelu", inputs=["x"], outputs=["y"])

x = np.array([-1, 0, 1]).astype(np.float32)
# expected output [-0.15865526, 0., 0.84134474]
y = (0.5 * x * (1 + np.vectorize(math.erf)(x / np.sqrt(2)))).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_gelu_default_1")

x = np.random.randn(3, 4, 5).astype(np.float32)
# expected output [2.99595031, 3.99987331, 4.99999857]
y = (0.5 * x * (1 + np.vectorize(math.erf)(x / np.sqrt(2)))).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_gelu_default_2")

gelu_tanh

node = onnx.helper.make_node(
    "Gelu", inputs=["x"], outputs=["y"], approximate="tanh"
)

x = np.array([-1, 0, 1]).astype(np.float32)
# expected output [-0.158808, 0., 0.841192]
y = (
    0.5
    * x
    * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * np.power(x, 3))))
).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_gelu_tanh_1")

x = np.random.randn(3, 4, 5).astype(np.float32)
# expected output [2.9963627, 3.99993, 4.9999995]
y = (
    0.5
    * x
    * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * np.power(x, 3))))
).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_gelu_tanh_2")

Gemm

There are 11 test cases, listed as following:

all_attributes

node = onnx.helper.make_node(
    "Gemm",
    inputs=["a", "b", "c"],
    outputs=["y"],
    alpha=0.25,
    beta=0.35,
    transA=1,
    transB=1,
)
a = np.random.ranf([4, 3]).astype(np.float32)
b = np.random.ranf([5, 4]).astype(np.float32)
c = np.random.ranf([1, 5]).astype(np.float32)
y = gemm_reference_implementation(
    a, b, c, transA=1, transB=1, alpha=0.25, beta=0.35
)
expect(node, inputs=[a, b, c], outputs=[y], name="test_gemm_all_attributes")

alpha

node = onnx.helper.make_node(
    "Gemm", inputs=["a", "b", "c"], outputs=["y"], alpha=0.5
)
a = np.random.ranf([3, 5]).astype(np.float32)
b = np.random.ranf([5, 4]).astype(np.float32)
c = np.zeros([1, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c, alpha=0.5)
expect(node, inputs=[a, b, c], outputs=[y], name="test_gemm_alpha")

beta

node = onnx.helper.make_node(
    "Gemm", inputs=["a", "b", "c"], outputs=["y"], beta=0.5
)
a = np.random.ranf([2, 7]).astype(np.float32)
b = np.random.ranf([7, 4]).astype(np.float32)
c = np.random.ranf([1, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c, beta=0.5)
expect(node, inputs=[a, b, c], outputs=[y], name="test_gemm_beta")

default_matrix_bias

node = onnx.helper.make_node("Gemm", inputs=["a", "b", "c"], outputs=["y"])
a = np.random.ranf([3, 6]).astype(np.float32)
b = np.random.ranf([6, 4]).astype(np.float32)
c = np.random.ranf([3, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c)
expect(
    node, inputs=[a, b, c], outputs=[y], name="test_gemm_default_matrix_bias"
)

default_no_bias

node = onnx.helper.make_node("Gemm", inputs=["a", "b"], outputs=["y"])
a = np.random.ranf([2, 10]).astype(np.float32)
b = np.random.ranf([10, 3]).astype(np.float32)
y = gemm_reference_implementation(a, b)
expect(node, inputs=[a, b], outputs=[y], name="test_gemm_default_no_bias")

default_scalar_bias

node = onnx.helper.make_node("Gemm", inputs=["a", "b", "c"], outputs=["y"])
a = np.random.ranf([2, 3]).astype(np.float32)
b = np.random.ranf([3, 4]).astype(np.float32)
c = np.array(3.14).astype(np.float32)
y = gemm_reference_implementation(a, b, c)
expect(
    node, inputs=[a, b, c], outputs=[y], name="test_gemm_default_scalar_bias"
)

default_single_elem_vector_bias

node = onnx.helper.make_node("Gemm", inputs=["a", "b", "c"], outputs=["y"])
a = np.random.ranf([3, 7]).astype(np.float32)
b = np.random.ranf([7, 3]).astype(np.float32)
c = np.random.ranf([1]).astype(np.float32)
y = gemm_reference_implementation(a, b, c)
expect(
    node,
    inputs=[a, b, c],
    outputs=[y],
    name="test_gemm_default_single_elem_vector_bias",
)

default_vector_bias

node = onnx.helper.make_node("Gemm", inputs=["a", "b", "c"], outputs=["y"])
a = np.random.ranf([2, 7]).astype(np.float32)
b = np.random.ranf([7, 4]).astype(np.float32)
c = np.random.ranf([1, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c)
expect(
    node, inputs=[a, b, c], outputs=[y], name="test_gemm_default_vector_bias"
)

default_zero_bias

node = onnx.helper.make_node("Gemm", inputs=["a", "b", "c"], outputs=["y"])
a = np.random.ranf([3, 5]).astype(np.float32)
b = np.random.ranf([5, 4]).astype(np.float32)
c = np.zeros([1, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c)
expect(node, inputs=[a, b, c], outputs=[y], name="test_gemm_default_zero_bias")

transposeA

node = onnx.helper.make_node(
    "Gemm", inputs=["a", "b", "c"], outputs=["y"], transA=1
)
a = np.random.ranf([6, 3]).astype(np.float32)
b = np.random.ranf([6, 4]).astype(np.float32)
c = np.zeros([1, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c, transA=1)
expect(node, inputs=[a, b, c], outputs=[y], name="test_gemm_transposeA")

transposeB

node = onnx.helper.make_node(
    "Gemm", inputs=["a", "b", "c"], outputs=["y"], transB=1
)
a = np.random.ranf([3, 6]).astype(np.float32)
b = np.random.ranf([4, 6]).astype(np.float32)
c = np.zeros([1, 4]).astype(np.float32)
y = gemm_reference_implementation(a, b, c, transB=1)
expect(node, inputs=[a, b, c], outputs=[y], name="test_gemm_transposeB")

GlobalAveragePool

There are 2 test cases, listed as following:

globalaveragepool

node = onnx.helper.make_node(
    "GlobalAveragePool",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(1, 3, 5, 5).astype(np.float32)
y = np.mean(x, axis=tuple(range(2, np.ndim(x))), keepdims=True)
expect(node, inputs=[x], outputs=[y], name="test_globalaveragepool")

globalaveragepool_precomputed

node = onnx.helper.make_node(
    "GlobalAveragePool",
    inputs=["x"],
    outputs=["y"],
)
x = np.array(
    [
        [
            [
                [1, 2, 3],
                [4, 5, 6],
                [7, 8, 9],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[5]]]]).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_globalaveragepool_precomputed")

GlobalMaxPool

There are 2 test cases, listed as following:

globalmaxpool

node = onnx.helper.make_node(
    "GlobalMaxPool",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(1, 3, 5, 5).astype(np.float32)
y = np.max(x, axis=tuple(range(2, np.ndim(x))), keepdims=True)
expect(node, inputs=[x], outputs=[y], name="test_globalmaxpool")

globalmaxpool_precomputed

node = onnx.helper.make_node(
    "GlobalMaxPool",
    inputs=["x"],
    outputs=["y"],
)
x = np.array(
    [
        [
            [
                [1, 2, 3],
                [4, 5, 6],
                [7, 8, 9],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[9]]]]).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_globalmaxpool_precomputed")

Gradient

There are 2 test cases, listed as following:

gradient_scalar_add

add_node = onnx.helper.make_node("Add", ["a", "b"], ["c"], name="my_add")
gradient_node = onnx.helper.make_node(
    "Gradient",
    ["a", "b"],
    ["dc_da", "dc_db"],
    name="my_gradient",
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
    xs=["a", "b"],
    y="c",
)

a = np.array(1.0).astype(np.float32)
b = np.array(2.0).astype(np.float32)
c = a + b
# dc / da = d(a+b) / da = 1
dc_da = np.array(1).astype(np.float32)
# db / db = d(a+b) / db = 1
dc_db = np.array(1).astype(np.float32)

graph = onnx.helper.make_graph(
    nodes=[add_node, gradient_node],
    name="GradientOfAdd",
    inputs=[
        onnx.helper.make_tensor_value_info("a", onnx.TensorProto.FLOAT, []),
        onnx.helper.make_tensor_value_info("b", onnx.TensorProto.FLOAT, []),
    ],
    outputs=[
        onnx.helper.make_tensor_value_info("c", onnx.TensorProto.FLOAT, []),
        onnx.helper.make_tensor_value_info("dc_da", onnx.TensorProto.FLOAT, []),
        onnx.helper.make_tensor_value_info("dc_db", onnx.TensorProto.FLOAT, []),
    ],
)
opsets = [
    onnx.helper.make_operatorsetid(ONNX_DOMAIN, 12),
    onnx.helper.make_operatorsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1),
]
model = onnx.helper.make_model_gen_version(
    graph, producer_name="backend-test", opset_imports=opsets
)
expect(
    model, inputs=[a, b], outputs=[c, dc_da, dc_db], name="test_gradient_of_add"
)

gradient_scalar_add_and_mul

add_node = onnx.helper.make_node("Add", ["a", "b"], ["c"], name="my_add")
mul_node = onnx.helper.make_node("Mul", ["c", "a"], ["d"], name="my_mul")
gradient_node = onnx.helper.make_node(
    "Gradient",
    ["a", "b"],
    ["dd_da", "dd_db"],
    name="my_gradient",
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
    xs=["a", "b"],
    y="d",
)

a = np.array(1.0).astype(np.float32)
b = np.array(2.0).astype(np.float32)
c = a + b
# d = a * c = a * (a + b)
d = a * c
# dd / da = d(a*a+a*b) / da = 2 * a + b
dd_da = (2 * a + b).astype(np.float32)
# dd / db = d(a*a+a*b) / db = a
dd_db = a

graph = onnx.helper.make_graph(
    nodes=[add_node, mul_node, gradient_node],
    name="GradientOfTwoOperators",
    inputs=[
        onnx.helper.make_tensor_value_info("a", onnx.TensorProto.FLOAT, []),
        onnx.helper.make_tensor_value_info("b", onnx.TensorProto.FLOAT, []),
    ],
    outputs=[
        onnx.helper.make_tensor_value_info("d", onnx.TensorProto.FLOAT, []),
        onnx.helper.make_tensor_value_info("dd_da", onnx.TensorProto.FLOAT, []),
        onnx.helper.make_tensor_value_info("dd_db", onnx.TensorProto.FLOAT, []),
    ],
)

opsets = [
    onnx.helper.make_operatorsetid(ONNX_DOMAIN, 12),
    onnx.helper.make_operatorsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1),
]
model = onnx.helper.make_model_gen_version(
    graph, producer_name="backend-test", opset_imports=opsets
)
expect(
    model,
    inputs=[a, b],
    outputs=[d, dd_da, dd_db],
    name="test_gradient_of_add_and_mul",
)

Greater

There are 4 test cases, listed as following:

greater

node = onnx.helper.make_node(
    "Greater",
    inputs=["x", "y"],
    outputs=["greater"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater")

x = np.random.randn(3, 4, 5).astype(np.int8)
y = np.random.randn(3, 4, 5).astype(np.int8)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_int8")

x = np.random.randn(3, 4, 5).astype(np.int16)
y = np.random.randn(3, 4, 5).astype(np.int16)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_int16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_uint64")

greater

node = onnx.helper.make_node(
    "GreaterOrEqual",
    inputs=["x", "y"],
    outputs=["greater_equal"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal")

x = np.random.randn(3, 4, 5).astype(np.int8)
y = np.random.randn(3, 4, 5).astype(np.int8)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_int8")

x = np.random.randn(3, 4, 5).astype(np.int16)
y = np.random.randn(3, 4, 5).astype(np.int16)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_int16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_uint64")

greater_broadcast

node = onnx.helper.make_node(
    "Greater",
    inputs=["x", "y"],
    outputs=["greater"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
z = np.greater(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_bcast")

greater_broadcast

node = onnx.helper.make_node(
    "GreaterOrEqual",
    inputs=["x", "y"],
    outputs=["greater_equal"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
z = np.greater_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_greater_equal_bcast")

GridSample

There are 4 test cases, listed as following:

gridsample

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
    padding_mode="zeros",
    align_corners=0,
)
# X shape, [N, C, H, W] - [1, 1, 4, 4]
X = np.array(
    [
        [
            [
                [0.0, 1.0, 2.0, 3.0],
                [4.0, 5.0, 6.0, 7.0],
                [8.0, 9.0, 10.0, 11.0],
                [12.0, 13.0, 14.0, 15.0],
            ]
        ]
    ],
    dtype=np.float32,
)
# Grid shape, [N, H_out, W_out, 2] - [1, 6, 6, 2]
Grid = np.array(
    [
        [
            [
                [-1.0000, -1.0000],
                [-0.6000, -1.0000],
                [-0.2000, -1.0000],
                [0.2000, -1.0000],
                [0.6000, -1.0000],
                [1.0000, -1.0000],
            ],
            [
                [-1.0000, -0.6000],
                [-0.6000, -0.6000],
                [-0.2000, -0.6000],
                [0.2000, -0.6000],
                [0.6000, -0.6000],
                [1.0000, -0.6000],
            ],
            [
                [-1.0000, -0.2000],
                [-0.6000, -0.2000],
                [-0.2000, -0.2000],
                [0.2000, -0.2000],
                [0.6000, -0.2000],
                [1.0000, -0.2000],
            ],
            [
                [-1.0000, 0.2000],
                [-0.6000, 0.2000],
                [-0.2000, 0.2000],
                [0.2000, 0.2000],
                [0.6000, 0.2000],
                [1.0000, 0.2000],
            ],
            [
                [-1.0000, 0.6000],
                [-0.6000, 0.6000],
                [-0.2000, 0.6000],
                [0.2000, 0.6000],
                [0.6000, 0.6000],
                [1.0000, 0.6000],
            ],
            [
                [-1.0000, 1.0000],
                [-0.6000, 1.0000],
                [-0.2000, 1.0000],
                [0.2000, 1.0000],
                [0.6000, 1.0000],
                [1.0000, 1.0000],
            ],
        ]
    ],
    dtype=np.float32,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 6, 6]
Y = np.array(
    [
        [
            [
                [0.0000, 0.1500, 0.5500, 0.9500, 1.3500, 0.7500],
                [0.6000, 1.5000, 2.3000, 3.1000, 3.9000, 2.1000],
                [2.2000, 4.7000, 5.5000, 6.3000, 7.1000, 3.7000],
                [3.8000, 7.9000, 8.7000, 9.5000, 10.3000, 5.3000],
                [5.4000, 11.1000, 11.9000, 12.7000, 13.5000, 6.9000],
                [3.0000, 6.1500, 6.5500, 6.9500, 7.3500, 3.7500],
            ]
        ]
    ],
    dtype=np.float32,
)
expect(node, inputs=[X, Grid], outputs=[Y], name="test_gridsample")

gridsample_mode_aligncorners

# X shape, [N, C, H, W] - [1, 1, 3, 2]
X = np.array(
    [[[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]]],
    dtype=np.float32,
)
# Grid shape, [N, H_out, W_out, 2] - [1, 2, 4, 2]
Grid = np.array(
    [
        [
            [
                [-1.0000, -1.0000],
                [-0.5000, -0.5000],
                [-0.2000, -0.2000],
                [0.0000, 0.0000],
            ],
            [
                [0.0000, 0.0000],
                [-0.2000, -0.2000],
                [0.5000, 0.5000],
                [1.0000, 1.0000],
            ],
        ]
    ],
    dtype=np.float32,
)

# setting mode = 'bilinear', default align_corners = 0
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bilinear = np.array(
    [[[[0.0000, 0.5000, 1.7000, 2.5000], [2.5000, 1.7000, 4.5000, 1.2500]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bilinear],
    name="test_gridsample_bilinear",
)

# setting mode = 'bilinear', align_corners = 1
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
    align_corners=1,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_align_corners = np.array(
    [[[[0.0000, 1.2500, 2.0000, 2.5000], [2.5000, 2.0000, 3.7500, 5.0000]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_align_corners],
    name="test_gridsample_aligncorners_true",
)

# setting mode = 'nearest'
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="nearest",
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_nearest = np.array(
    [[[[0.0, 0.0, 2.0, 2.0], [2.0, 2.0, 5.0, 0.0]]]],
    dtype=np.float32,
)

expect(
    node, inputs=[X, Grid], outputs=[Y_nearest], name="test_gridsample_nearest"
)

# setting mode = 'bicubic'
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="cubic",
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bicubic = np.array(
    [[[[-0.1406, 0.3828, 1.7556, 2.9688], [2.9688, 1.7556, 5.1445, 1.3906]]]],
    dtype=np.float32,
)

expect(
    node, inputs=[X, Grid], outputs=[Y_bicubic], name="test_gridsample_bicubic"
)

# ============================================================================
# Additional tests
# The reference output tensors were generated using PyTorch 2.0.
Grid = np.array(
    [
        [
            [[-1.0, -0.8], [-0.6, -0.5], [-0.1, -0.2], [0.7, 0.0]],
            [[0.0, 0.4], [0.2, -0.2], [-0.3, 0.5], [-1.0, 1.0]],
        ]
    ],
    dtype=np.float32,
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="nearest",
    align_corners=0,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_nearest = np.array(
    [[[[0.0, 0.0, 2.0, 3.0], [4.0, 3.0, 4.0, 4.0]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_nearest],
    name="test_gridsample_nearest_align_corners_0_additional_1",
)

# setting mode = 'nearest'
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="nearest",
    align_corners=1,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_nearest = np.array(
    [[[[0.0, 0.0, 2.0, 3.0], [2.0, 3.0, 4.0, 4.0]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_nearest],
    name="test_gridsample_nearest_align_corners_1_additional_1",
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
    align_corners=0,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bilinear = np.array(
    [[[[0.0000, 0.4500, 1.8000, 2.4000], [3.7000, 2.1000, 3.7000, 1.0000]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bilinear],
    name="test_gridsample_bilinear_align_corners_0_additional_1",
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
    align_corners=1,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bilinear = np.array(
    [[[[0.4000, 1.2000, 2.0500, 2.8500], [3.3000, 2.2000, 3.3500, 4.0000]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bilinear],
    name="test_gridsample_bilinear_align_corners_1_additional_1",
)

# These two new bicubic tests produces slightly higher error ~5e-5
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="cubic",
    align_corners=0,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bicubic = np.array(
    [
        [
            [
                [-0.173250, 0.284265, 1.923106, 2.568000],
                [5.170375, 2.284414, 4.744844, 1.046875],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bicubic],
    name="test_gridsample_bicubic_align_corners_0_additional_1",
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="cubic",
    align_corners=1,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bicubic = np.array(
    [
        [
            [
                [0.304001, 1.128750, 2.266270, 3.144844],
                [4.531500, 2.455360, 4.599819, 4.000000],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bicubic],
    name="test_gridsample_bicubic_align_corners_1_additional_1",
)

gridsample_paddingmode

# X shape, [N, C, H, W] - [1, 1, 3, 2]
X = np.array(
    [[[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]]],
    dtype=np.float32,
)
# Grid shape, [N, H_out, W_out, 2] - [1, 2, 4, 2]
Grid = np.array(
    [
        [
            [
                [-10.0000, -10.0000],
                [-5.0000, -5.0000],
                [-0.2000, -0.2000],
                [10.0000, 10.0000],
            ],
            [
                [10.0000, 10.0000],
                [-0.2000, -0.2000],
                [5.0000, 5.0000],
                [10.0000, 10.0000],
            ],
        ]
    ],
    dtype=np.float32,
)

# setting padding_mode = 'zeros'
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    padding_mode="zeros",
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_zeros = np.array(
    [[[[0.0000, 0.0000, 1.7000, 0.0000], [0.0000, 1.7000, 0.0000, 0.0000]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_zeros],
    name="test_gridsample_zeros_padding",
)

# setting padding_mode = 'border'
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    padding_mode="border",
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_border = np.array(
    [[[[0.0000, 0.0000, 1.7000, 5.0000], [5.0000, 1.7000, 5.0000, 5.0000]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_border],
    name="test_gridsample_border_padding",
)

# setting padding_mode = 'reflection'
node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    padding_mode="reflection",
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_reflection = np.array(
    [[[[2.5000, 0.0000, 1.7000, 2.5000], [2.5000, 1.7000, 5.0000, 2.5000]]]],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_reflection],
    name="test_gridsample_reflection_padding",
)

volumeetric_gridsample_mode_aligncorners

X = np.array(
    [
        [
            [
                [[1.0, 2.0], [3.0, 4.0]],
                [[5.0, 6.0], [7.0, 8.0]],
                [[9.0, 10.0], [11.0, 12.0]],
            ]
        ]
    ],
    dtype=np.float32,
)

Grid = np.array(
    [
        [
            [
                [[-1.0, -1.0, -1.0], [-1.0, -0.5, 0.3]],
                [[-0.5, -0.5, -0.5], [1.0, -0.6, -1.0]],
                [[-0.2, -0.2, -0.2], [0.4, 0.2, 0.6]],
                [[0.0, 0.0, 0.0], [-1.0, 0.0, 0.0]],
            ],
            [
                [[0.0, 0.0, 0.0], [-1.0, 1.0, 0.0]],
                [[-0.2, -0.2, -0.2], [1.0, 0.4, -0.2]],
                [[0.5, 0.5, 0.5], [-1.0, -0.8, 0.8]],
                [[1.0, 1.0, 1.0], [0.4, 0.6, -0.3]],
            ],
        ]
    ],
    dtype=np.float32,
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="nearest",
    align_corners=0,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_nearest = np.array(
    [
        [
            [
                [[1.0, 5.0], [1.0, 0.0], [5.0, 12.0], [5.0, 5.0]],
                [[5.0, 0.0], [5.0, 0.0], [12.0, 9.0], [0.0, 8.0]],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_nearest],
    name="test_gridsample_volumetric_nearest_align_corners_0",
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="nearest",
    align_corners=1,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_nearest = np.array(
    [
        [
            [
                [[1.0, 5.0], [1.0, 2.0], [5.0, 12.0], [5.0, 5.0]],
                [[5.0, 7.0], [5.0, 8.0], [12.0, 9.0], [12.0, 8.0]],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_nearest],
    name="test_gridsample_volumetric_nearest_align_corners_1",
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
    align_corners=0,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bilinear = np.array(
    [
        [
            [
                [
                    [0.1250, 3.4000],
                    [2.0000, 0.4500],
                    [4.7000, 10.9000],
                    [6.5000, 3.0000],
                ],
                [
                    [6.5000, 1.7500],
                    [4.7000, 3.3000],
                    [11.0000, 2.5200],
                    [1.5000, 5.4900],
                ],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bilinear],
    name="test_gridsample_volumetric_bilinear_align_corners_0",
)

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="linear",
    align_corners=1,
)
# Y shape, [N, C, H_out, W_out] - [1, 1, 2, 4]
Y_bilinear = np.array(
    [
        [
            [
                [
                    [1.0000, 6.7000],
                    [3.7500, 2.4000],
                    [5.4000, 9.3000],
                    [6.5000, 6.0000],
                ],
                [
                    [6.5000, 7.0000],
                    [5.4000, 6.6000],
                    [9.2500, 8.4000],
                    [12.0000, 6.1000],
                ],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, Grid],
    outputs=[Y_bilinear],
    name="test_gridsample_volumetric_bilinear_align_corners_1",
)

GroupNormalization

There are 2 test cases, listed as following:

epsilon

c = 4
num_groups = 2
x = np.random.randn(3, c, 2, 2).astype(np.float32)
scale = np.random.randn(c).astype(np.float32)
bias = np.random.randn(c).astype(np.float32)
epsilon = 1e-2
y = _group_normalization(x, num_groups, scale, bias, epsilon).astype(np.float32)

node = onnx.helper.make_node(
    "GroupNormalization",
    inputs=["x", "scale", "bias"],
    outputs=["y"],
    epsilon=epsilon,
    num_groups=num_groups,
)

expect(
    node,
    inputs=[x, scale, bias],
    outputs=[y],
    name="test_group_normalization_epsilon",
)

groupnormalization

c = 4
num_groups = 2
x = np.random.randn(3, c, 2, 2).astype(np.float32)
scale = np.random.randn(c).astype(np.float32)
bias = np.random.randn(c).astype(np.float32)
y = _group_normalization(x, num_groups, scale, bias).astype(np.float32)

node = onnx.helper.make_node(
    "GroupNormalization",
    inputs=["x", "scale", "bias"],
    outputs=["y"],
    num_groups=num_groups,
)

expect(
    node,
    inputs=[x, scale, bias],
    outputs=[y],
    name="test_group_normalization_example",
)

HammingWindow

There are 1 test cases, listed as following:

hammingwindow

# Test periodic window
node = onnx.helper.make_node(
    "HammingWindow",
    inputs=["x"],
    outputs=["y"],
)
size = np.int32(10)
a0 = 25 / 46
a1 = 1 - a0
y = a0 - a1 * np.cos(2 * np.pi * np.arange(0, size, 1, dtype=np.float32) / size)
expect(
    node,
    inputs=[size],
    outputs=[y.astype(np.float32)],
    name="test_hammingwindow",
)

# Test symmetric window
node = onnx.helper.make_node(
    "HammingWindow", inputs=["x"], outputs=["y"], periodic=0
)
size = np.int32(10)
a0 = 25 / 46
a1 = 1 - a0
y = a0 - a1 * np.cos(
    2 * np.pi * np.arange(0, size, 1, dtype=np.float32) / (size - 1)
)
expect(
    node,
    inputs=[size],
    outputs=[y.astype(np.float32)],
    name="test_hammingwindow_symmetric",
)

HannWindow

There are 1 test cases, listed as following:

hannwindow

# Test periodic window
node = onnx.helper.make_node(
    "HannWindow",
    inputs=["x"],
    outputs=["y"],
)
size = np.int32(10)
a0 = 0.5
a1 = 0.5
y = a0 - a1 * np.cos(2 * np.pi * np.arange(0, size, 1, dtype=np.float32) / size)
expect(
    node, inputs=[size], outputs=[y.astype(np.float32)], name="test_hannwindow"
)

# Test symmetric window
node = onnx.helper.make_node(
    "HannWindow", inputs=["x"], outputs=["y"], periodic=0
)
size = np.int32(10)
a0 = 0.5
a1 = 0.5
y = a0 - a1 * np.cos(
    2 * np.pi * np.arange(0, size, 1, dtype=np.float32) / (size - 1)
)
expect(
    node,
    inputs=[size],
    outputs=[y.astype(np.float32)],
    name="test_hannwindow_symmetric",
)

HardSigmoid

There are 2 test cases, listed as following:

hardsigmoid

node = onnx.helper.make_node(
    "HardSigmoid", inputs=["x"], outputs=["y"], alpha=0.5, beta=0.6
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.clip(x * 0.5 + 0.6, 0, 1)  # expected output [0.1, 0.6, 1.]
expect(node, inputs=[x], outputs=[y], name="test_hardsigmoid_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x * 0.5 + 0.6, 0, 1)
expect(node, inputs=[x], outputs=[y], name="test_hardsigmoid")

hardsigmoid_default

default_alpha = 0.2
default_beta = 0.5
node = onnx.helper.make_node(
    "HardSigmoid",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x * default_alpha + default_beta, 0, 1)
expect(node, inputs=[x], outputs=[y], name="test_hardsigmoid_default")

HardSwish

There are 1 test cases, listed as following:

hardswish

node = onnx.helper.make_node(
    "HardSwish",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = hardswish(x)

expect(node, inputs=[x], outputs=[y], name="test_hardswish")

Hardmax

There are 2 test cases, listed as following:

hardmax

node = onnx.helper.make_node(
    "Hardmax",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([[3, 0, 1, 2], [2, 5, 1, 0], [0, 1, 3, 2], [0, 1, 2, 3]]).astype(
    np.float32
)
# expect result:
# [[1. 0. 0. 0.]
# [0. 1. 0. 0.]
# [0. 0. 1. 0.]
# [0. 0. 0. 1.]]
y = hardmax(x)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_example")

# For multiple occurrences of the maximal values, the first occurrence is selected for one-hot output
x = np.array([[3, 3, 3, 1]]).astype(np.float32)
# expect result:
# [[1, 0, 0, 0]]
y = hardmax(x)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_one_hot")

hardmax_axis

x = np.random.randn(3, 4, 5).astype(np.float32)
node = onnx.helper.make_node(
    "Hardmax",
    inputs=["x"],
    outputs=["y"],
    axis=0,
)
y = hardmax(x, axis=0)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_axis_0")

node = onnx.helper.make_node(
    "Hardmax",
    inputs=["x"],
    outputs=["y"],
    axis=1,
)
y = hardmax(x, axis=1)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_axis_1")

node = onnx.helper.make_node(
    "Hardmax",
    inputs=["x"],
    outputs=["y"],
    axis=2,
)
y = hardmax(x, axis=2)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_axis_2")

node = onnx.helper.make_node(
    "Hardmax",
    inputs=["x"],
    outputs=["y"],
    axis=-1,
)
y = hardmax(x, axis=-1)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_negative_axis")

# default axis is -1
node = onnx.helper.make_node(
    "Hardmax",
    inputs=["x"],
    outputs=["y"],
)
expect(node, inputs=[x], outputs=[y], name="test_hardmax_default_axis")

Identity

There are 3 test cases, listed as following:

identity

node = onnx.helper.make_node(
    "Identity",
    inputs=["x"],
    outputs=["y"],
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(node, inputs=[data], outputs=[data], name="test_identity")

identity_opt

ten_in_tp = onnx.helper.make_tensor_type_proto(
    onnx.TensorProto.FLOAT, shape=[5]
)
seq_in_tp = onnx.helper.make_sequence_type_proto(ten_in_tp)
opt_in_tp = onnx.helper.make_optional_type_proto(seq_in_tp)

identity_node = onnx.helper.make_node(
    "Identity", inputs=["opt_in"], outputs=["opt_out"]
)

x = [np.array([1, 2, 3, 4, 5]).astype(np.float32)]

expect(
    identity_node,
    inputs=[x],
    outputs=[x],
    name="test_identity_opt",
    opset_imports=[onnx.helper.make_opsetid("", 16)],
    input_type_protos=[opt_in_tp],
    output_type_protos=[opt_in_tp],
)

sequence

node = onnx.helper.make_node(
    "Identity",
    inputs=["x"],
    outputs=["y"],
)

data = [
    np.array(
        [
            [
                [
                    [1, 2],
                    [3, 4],
                ]
            ]
        ],
        dtype=np.float32,
    ),
    np.array(
        [
            [
                [
                    [2, 3],
                    [1, 5],
                ]
            ]
        ],
        dtype=np.float32,
    ),
]

expect(node, inputs=[data], outputs=[data], name="test_identity_sequence")

If

There are 3 test cases, listed as following:

if

# Given a bool scalar input cond.
# return constant tensor x if cond is True, otherwise return constant tensor y.

then_out = onnx.helper.make_tensor_value_info(
    "then_out", onnx.TensorProto.FLOAT, [5]
)
else_out = onnx.helper.make_tensor_value_info(
    "else_out", onnx.TensorProto.FLOAT, [5]
)

x = np.array([1, 2, 3, 4, 5]).astype(np.float32)
y = np.array([5, 4, 3, 2, 1]).astype(np.float32)

then_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["then_out"],
    value=onnx.numpy_helper.from_array(x),
)

else_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["else_out"],
    value=onnx.numpy_helper.from_array(y),
)

then_body = onnx.helper.make_graph(
    [then_const_node], "then_body", [], [then_out]
)

else_body = onnx.helper.make_graph(
    [else_const_node], "else_body", [], [else_out]
)

if_node = onnx.helper.make_node(
    "If",
    inputs=["cond"],
    outputs=["res"],
    then_branch=then_body,
    else_branch=else_body,
)

cond = np.array(1).astype(bool)
res = x if cond else y
expect(
    if_node,
    inputs=[cond],
    outputs=[res],
    name="test_if",
    opset_imports=[onnx.helper.make_opsetid("", 11)],
)

if_optional

# Given a bool scalar input cond, return an empty optional sequence of
# tensor if True, return an optional sequence with value x
# (the input optional sequence) otherwise.

ten_in_tp = onnx.helper.make_tensor_type_proto(
    onnx.TensorProto.FLOAT, shape=[5]
)
seq_in_tp = onnx.helper.make_sequence_type_proto(ten_in_tp)

then_out_tensor_tp = onnx.helper.make_tensor_type_proto(
    onnx.TensorProto.FLOAT, shape=[5]
)
then_out_seq_tp = onnx.helper.make_sequence_type_proto(then_out_tensor_tp)
then_out_opt_tp = onnx.helper.make_optional_type_proto(then_out_seq_tp)
then_out = onnx.helper.make_value_info("optional_empty", then_out_opt_tp)

else_out_tensor_tp = onnx.helper.make_tensor_type_proto(
    onnx.TensorProto.FLOAT, shape=[5]
)
else_out_seq_tp = onnx.helper.make_sequence_type_proto(else_out_tensor_tp)
else_out_opt_tp = onnx.helper.make_optional_type_proto(else_out_seq_tp)
else_out = onnx.helper.make_value_info("else_opt", else_out_opt_tp)

x = [np.array([1, 2, 3, 4, 5]).astype(np.float32)]
cond = np.array(0).astype(bool)
res = compute_if_outputs(x, cond)

opt_empty_in = onnx.helper.make_node(
    "Optional", inputs=[], outputs=["optional_empty"], type=seq_in_tp
)

then_body = onnx.helper.make_graph([opt_empty_in], "then_body", [], [then_out])

else_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["x"],
    value=onnx.numpy_helper.from_array(x[0]),
)

else_seq_node = onnx.helper.make_node(
    "SequenceConstruct", inputs=["x"], outputs=["else_seq"]
)

else_optional_seq_node = onnx.helper.make_node(
    "Optional", inputs=["else_seq"], outputs=["else_opt"]
)

else_body = onnx.helper.make_graph(
    [else_const_node, else_seq_node, else_optional_seq_node],
    "else_body",
    [],
    [else_out],
)

if_node = onnx.helper.make_node(
    "If",
    inputs=["cond"],
    outputs=["sequence"],
    then_branch=then_body,
    else_branch=else_body,
)

expect(
    if_node,
    inputs=[cond],
    outputs=[res],
    name="test_if_opt",
    output_type_protos=[else_out_opt_tp],
    opset_imports=[onnx.helper.make_opsetid("", 16)],
)

if_seq

# Given a bool scalar input cond.
# return constant sequence x if cond is True, otherwise return constant sequence y.

then_out = onnx.helper.make_tensor_sequence_value_info(
    "then_out", onnx.TensorProto.FLOAT, shape=[5]
)
else_out = onnx.helper.make_tensor_sequence_value_info(
    "else_out", onnx.TensorProto.FLOAT, shape=[5]
)

x = [np.array([1, 2, 3, 4, 5]).astype(np.float32)]
y = [np.array([5, 4, 3, 2, 1]).astype(np.float32)]

then_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["x"],
    value=onnx.numpy_helper.from_array(x[0]),
)

then_seq_node = onnx.helper.make_node(
    "SequenceConstruct", inputs=["x"], outputs=["then_out"]
)

else_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["y"],
    value=onnx.numpy_helper.from_array(y[0]),
)

else_seq_node = onnx.helper.make_node(
    "SequenceConstruct", inputs=["y"], outputs=["else_out"]
)

then_body = onnx.helper.make_graph(
    [then_const_node, then_seq_node], "then_body", [], [then_out]
)

else_body = onnx.helper.make_graph(
    [else_const_node, else_seq_node], "else_body", [], [else_out]
)

if_node = onnx.helper.make_node(
    "If",
    inputs=["cond"],
    outputs=["res"],
    then_branch=then_body,
    else_branch=else_body,
)

cond = np.array(1).astype(bool)
res = x if cond else y
expect(
    if_node,
    inputs=[cond],
    outputs=[res],
    name="test_if_seq",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

ImageDecoder

There are 9 test cases, listed as following:

image_decoder_decode_bmp_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "bmp", _image_decoder_data.image_decoder_decode_bmp_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_bmp_rgb",
)

image_decoder_decode_jpeg2k_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "jpeg2000", _image_decoder_data.image_decoder_decode_jpeg2k_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_jpeg2k_rgb",
)

image_decoder_decode_jpeg_bgr

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="BGR",
)

data, output = _generate_test_data(
    "jpeg", _image_decoder_data.image_decoder_decode_jpeg_bgr, "BGR"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_jpeg_bgr",
)

image_decoder_decode_jpeg_grayscale

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="Grayscale",
)

data, output = _generate_test_data(
    "jpeg", _image_decoder_data.image_decoder_decode_jpeg_grayscale, "Grayscale"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_jpeg_grayscale",
)

image_decoder_decode_jpeg_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "jpeg", _image_decoder_data.image_decoder_decode_jpeg_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_jpeg_rgb",
)

image_decoder_decode_png_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "png", _image_decoder_data.image_decoder_decode_png_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_png_rgb",
)

image_decoder_decode_pnm_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "ppm", _image_decoder_data.image_decoder_decode_pnm_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_pnm_rgb",
)

image_decoder_decode_tiff_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "tiff", _image_decoder_data.image_decoder_decode_tiff_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_tiff_rgb",
)

image_decoder_decode_webp_rgb

node = onnx.helper.make_node(
    "ImageDecoder",
    inputs=["data"],
    outputs=["output"],
    pixel_format="RGB",
)

data, output = _generate_test_data(
    "webp", _image_decoder_data.image_decoder_decode_webp_rgb, "RGB"
)
expect(
    node,
    inputs=[data],
    outputs=[output],
    name="test_image_decoder_decode_webp_rgb",
)

InstanceNormalization

There are 1 test cases, listed as following:

instancenormalization

def _instancenorm_test_mode(
    x: np.ndarray, s: np.ndarray, bias: np.ndarray, epsilon: float = 1e-5
) -> np.ndarray:
    dims_x = len(x.shape)
    axis = tuple(range(2, dims_x))
    mean = np.mean(x, axis=axis, keepdims=True)
    var = np.var(x, axis=axis, keepdims=True)
    dim_ones = (1,) * (dims_x - 2)
    s = s.reshape(-1, *dim_ones)
    bias = bias.reshape(-1, *dim_ones)
    return s * (x - mean) / np.sqrt(var + epsilon) + bias

# input size: (1, 2, 1, 3)
x = np.array([[[[-1, 0, 1]], [[2, 3, 4]]]]).astype(np.float32)
s = np.array([1.0, 1.5]).astype(np.float32)
bias = np.array([0, 1]).astype(np.float32)
y = _instancenorm_test_mode(x, s, bias).astype(np.float32)

node = onnx.helper.make_node(
    "InstanceNormalization",
    inputs=["x", "s", "bias"],
    outputs=["y"],
)

# output size: (1, 2, 1, 3)
expect(node, inputs=[x, s, bias], outputs=[y], name="test_instancenorm_example")

# input size: (2, 3, 4, 5)
x = np.random.randn(2, 3, 4, 5).astype(np.float32)
s = np.random.randn(3).astype(np.float32)
bias = np.random.randn(3).astype(np.float32)
epsilon = 1e-2
y = _instancenorm_test_mode(x, s, bias, epsilon).astype(np.float32)

node = onnx.helper.make_node(
    "InstanceNormalization",
    inputs=["x", "s", "bias"],
    outputs=["y"],
    epsilon=epsilon,
)

# output size: (2, 3, 4, 5)
expect(node, inputs=[x, s, bias], outputs=[y], name="test_instancenorm_epsilon")

IsInf

There are 4 test cases, listed as following:

infinity

node = onnx.helper.make_node(
    "IsInf",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1.2, np.nan, np.inf, 2.8, -np.inf, np.inf], dtype=np.float32)
y = np.isinf(x)
expect(node, inputs=[x], outputs=[y], name="test_isinf")

infinity_float16

node = onnx.helper.make_node(
    "IsInf",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1.2, np.nan, np.inf, 2.8, -np.inf, np.inf], dtype=np.float16)
y = np.isinf(x)
expect(node, inputs=[x], outputs=[y], name="test_isinf_float16")

negative_infinity_only

node = onnx.helper.make_node(
    "IsInf", inputs=["x"], outputs=["y"], detect_positive=0
)

x = np.array([-1.7, np.nan, np.inf, -3.6, -np.inf, np.inf], dtype=np.float32)
y = np.isneginf(x)
expect(node, inputs=[x], outputs=[y], name="test_isinf_negative")

positive_infinity_only

node = onnx.helper.make_node(
    "IsInf", inputs=["x"], outputs=["y"], detect_negative=0
)

x = np.array([-1.7, np.nan, np.inf, 3.6, -np.inf, np.inf], dtype=np.float32)
y = np.isposinf(x)
expect(node, inputs=[x], outputs=[y], name="test_isinf_positive")

IsNaN

There are 2 test cases, listed as following:

float16

node = onnx.helper.make_node(
    "IsNaN",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1.2, np.nan, np.inf, 2.8, -np.inf, np.inf], dtype=np.float16)
y = np.isnan(x)
expect(node, inputs=[x], outputs=[y], name="test_isnan_float16")

isnan

node = onnx.helper.make_node(
    "IsNaN",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1.2, np.nan, np.inf, 2.8, -np.inf, np.inf], dtype=np.float32)
y = np.isnan(x)
expect(node, inputs=[x], outputs=[y], name="test_isnan")

LRN

There are 2 test cases, listed as following:

default

alpha = 0.0001
beta = 0.75
bias = 1.0
nsize = 3
node = onnx.helper.make_node("LRN", inputs=["x"], outputs=["y"], size=3)
x = np.random.randn(5, 5, 5, 5).astype(np.float32)
square_sum = np.zeros((5, 5, 5, 5)).astype(np.float32)
for n, c, h, w in np.ndindex(x.shape):
    square_sum[n, c, h, w] = sum(
        x[
            n,
            max(0, c - math.floor((nsize - 1) / 2)) : min(
                5, c + math.ceil((nsize - 1) / 2) + 1
            ),
            h,
            w,
        ]
        ** 2
    )
y = x / ((bias + (alpha / nsize) * square_sum) ** beta)
expect(node, inputs=[x], outputs=[y], name="test_lrn_default")

lrn

alpha = 0.0002
beta = 0.5
bias = 2.0
nsize = 3
node = onnx.helper.make_node(
    "LRN",
    inputs=["x"],
    outputs=["y"],
    alpha=alpha,
    beta=beta,
    bias=bias,
    size=nsize,
)
x = np.random.randn(5, 5, 5, 5).astype(np.float32)
square_sum = np.zeros((5, 5, 5, 5)).astype(np.float32)
for n, c, h, w in np.ndindex(x.shape):
    square_sum[n, c, h, w] = sum(
        x[
            n,
            max(0, c - math.floor((nsize - 1) / 2)) : min(
                5, c + math.ceil((nsize - 1) / 2) + 1
            ),
            h,
            w,
        ]
        ** 2
    )
y = x / ((bias + (alpha / nsize) * square_sum) ** beta)
expect(node, inputs=[x], outputs=[y], name="test_lrn")

LSTM

There are 4 test cases, listed as following:

batchwise

input = np.array([[[1.0, 2.0]], [[3.0, 4.0]], [[5.0, 6.0]]]).astype(np.float32)

input_size = 2
hidden_size = 7
weight_scale = 0.3
number_of_gates = 4
layout = 1

node = onnx.helper.make_node(
    "LSTM",
    inputs=["X", "W", "R"],
    outputs=["Y", "Y_h"],
    hidden_size=hidden_size,
    layout=layout,
)

W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)

lstm = LSTMHelper(X=input, W=W, R=R, layout=layout)
Y, Y_h = lstm.step()
expect(
    node,
    inputs=[input, W, R],
    outputs=[Y.astype(np.float32), Y_h.astype(np.float32)],
    name="test_lstm_batchwise",
)

defaults

input = np.array([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]).astype(np.float32)

input_size = 2
hidden_size = 3
weight_scale = 0.1
number_of_gates = 4

node = onnx.helper.make_node(
    "LSTM", inputs=["X", "W", "R"], outputs=["", "Y_h"], hidden_size=hidden_size
)

W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)

lstm = LSTMHelper(X=input, W=W, R=R)
_, Y_h = lstm.step()
expect(
    node,
    inputs=[input, W, R],
    outputs=[Y_h.astype(np.float32)],
    name="test_lstm_defaults",
)

initial_bias

input = np.array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]).astype(
    np.float32
)

input_size = 3
hidden_size = 4
weight_scale = 0.1
custom_bias = 0.1
number_of_gates = 4

node = onnx.helper.make_node(
    "LSTM",
    inputs=["X", "W", "R", "B"],
    outputs=["", "Y_h"],
    hidden_size=hidden_size,
)

W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)

# Adding custom bias
W_B = custom_bias * np.ones((1, number_of_gates * hidden_size)).astype(
    np.float32
)
R_B = np.zeros((1, number_of_gates * hidden_size)).astype(np.float32)
B = np.concatenate((W_B, R_B), 1)

lstm = LSTMHelper(X=input, W=W, R=R, B=B)
_, Y_h = lstm.step()
expect(
    node,
    inputs=[input, W, R, B],
    outputs=[Y_h.astype(np.float32)],
    name="test_lstm_with_initial_bias",
)

peepholes

input = np.array([[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]).astype(
    np.float32
)

input_size = 4
hidden_size = 3
weight_scale = 0.1
number_of_gates = 4
number_of_peepholes = 3

node = onnx.helper.make_node(
    "LSTM",
    inputs=["X", "W", "R", "B", "sequence_lens", "initial_h", "initial_c", "P"],
    outputs=["", "Y_h"],
    hidden_size=hidden_size,
)

# Initializing Inputs
W = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, input_size)
).astype(np.float32)
R = weight_scale * np.ones(
    (1, number_of_gates * hidden_size, hidden_size)
).astype(np.float32)
B = np.zeros((1, 2 * number_of_gates * hidden_size)).astype(np.float32)
seq_lens = np.repeat(input.shape[0], input.shape[1]).astype(np.int32)
init_h = np.zeros((1, input.shape[1], hidden_size)).astype(np.float32)
init_c = np.zeros((1, input.shape[1], hidden_size)).astype(np.float32)
P = weight_scale * np.ones((1, number_of_peepholes * hidden_size)).astype(
    np.float32
)

lstm = LSTMHelper(
    X=input, W=W, R=R, B=B, P=P, initial_c=init_c, initial_h=init_h
)
_, Y_h = lstm.step()
expect(
    node,
    inputs=[input, W, R, B, seq_lens, init_h, init_c, P],
    outputs=[Y_h.astype(np.float32)],
    name="test_lstm_with_peepholes",
)

LayerNormalization

There are 4 test cases, listed as following:

d

X = np.random.randn(3, 4).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    B = np.random.randn(*normalized_shape).astype(np.float32)
    Y, mean, inv_std_dev = _layer_normalization(X, W, B, axis=axis)

    node = onnx.helper.make_node(
        "LayerNormalization",
        inputs=["X", "W", "B"],
        outputs=["Y", "Mean", "InvStdDev"],
        axis=axis,
    )

    if axis < 0:
        name = f"test_layer_normalization_2d_axis_negative_{-axis}"
    else:
        name = f"test_layer_normalization_2d_axis{axis}"

    expect(node, inputs=[X, W, B], outputs=[Y, mean, inv_std_dev], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

d_epsilon

epsilon = 1e-1
X = np.random.randn(2, 3, 5).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    B = np.random.randn(*normalized_shape).astype(np.float32)
    Y, mean, inv_std_dev = _layer_normalization(X, W, B, axis, epsilon)
    node = onnx.helper.make_node(
        "LayerNormalization",
        inputs=["X", "W", "B"],
        outputs=["Y", "Mean", "InvStdDev"],
        axis=axis,
        epsilon=epsilon,
    )

    if axis < 0:
        name = f"test_layer_normalization_3d_axis_negative_{-axis}_epsilon"
    else:
        name = f"test_layer_normalization_3d_axis{axis}_epsilon"

    expect(node, inputs=[X, W, B], outputs=[Y, mean, inv_std_dev], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

default_axis

X = np.random.randn(2, 3, 4, 5).astype(np.float32)

# Default axis in LayerNormalization is -1.
normalized_shape = calculate_normalized_shape(X.shape, -1)
W = np.random.randn(*normalized_shape).astype(np.float32)
B = np.random.randn(*normalized_shape).astype(np.float32)
# Axis is default to -1 in the reference implementation.
Y, mean, inv_std_dev = _layer_normalization(X, W, B)

# Not specifying axis attribute means -1.
node = onnx.helper.make_node(
    "LayerNormalization",
    inputs=["X", "W", "B"],
    outputs=["Y", "Mean", "InvStdDev"],
)

expect(
    node,
    inputs=[X, W, B],
    outputs=[Y, mean, inv_std_dev],
    name="test_layer_normalization_default_axis",
)

layernormalization

X = np.random.randn(2, 3, 4, 5).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    B = np.random.randn(*normalized_shape).astype(np.float32)
    Y, mean, inv_std_dev = _layer_normalization(X, W, B, axis)

    node = onnx.helper.make_node(
        "LayerNormalization",
        inputs=["X", "W", "B"],
        outputs=["Y", "Mean", "InvStdDev"],
        axis=axis,
    )

    if axis < 0:
        name = f"test_layer_normalization_4d_axis_negative_{-axis}"
    else:
        name = f"test_layer_normalization_4d_axis{axis}"

    expect(node, inputs=[X, W, B], outputs=[Y, mean, inv_std_dev], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

LeakyRelu

There are 2 test cases, listed as following:

leakyrelu

node = onnx.helper.make_node(
    "LeakyRelu", inputs=["x"], outputs=["y"], alpha=0.1
)

x = np.array([-1, 0, 1]).astype(np.float32)
# expected output [-0.1, 0., 1.]
y = np.clip(x, 0, np.inf) + np.clip(x, -np.inf, 0) * 0.1
expect(node, inputs=[x], outputs=[y], name="test_leakyrelu_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, 0, np.inf) + np.clip(x, -np.inf, 0) * 0.1
expect(node, inputs=[x], outputs=[y], name="test_leakyrelu")

leakyrelu_default

default_alpha = 0.01
node = onnx.helper.make_node(
    "LeakyRelu",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, 0, np.inf) + np.clip(x, -np.inf, 0) * default_alpha
expect(node, inputs=[x], outputs=[y], name="test_leakyrelu_default")

Less

There are 4 test cases, listed as following:

less

node = onnx.helper.make_node(
    "Less",
    inputs=["x", "y"],
    outputs=["less"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less")

x = np.random.randn(3, 4, 5).astype(np.int8)
y = np.random.randn(3, 4, 5).astype(np.int8)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_int8")

x = np.random.randn(3, 4, 5).astype(np.int16)
y = np.random.randn(3, 4, 5).astype(np.int16)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_int16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_uint64")

less

node = onnx.helper.make_node(
    "LessOrEqual",
    inputs=["x", "y"],
    outputs=["less_equal"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal")

x = np.random.randn(3, 4, 5).astype(np.int8)
y = np.random.randn(3, 4, 5).astype(np.int8)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_int8")

x = np.random.randn(3, 4, 5).astype(np.int16)
y = np.random.randn(3, 4, 5).astype(np.int16)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_int16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_uint8")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_uint16")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_uint32")

x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_uint64")

less_broadcast

node = onnx.helper.make_node(
    "Less",
    inputs=["x", "y"],
    outputs=["less"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
z = np.less(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_bcast")

less_broadcast

node = onnx.helper.make_node(
    "LessOrEqual",
    inputs=["x", "y"],
    outputs=["less_equal"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
z = np.less_equal(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_less_equal_bcast")

LinearAttention

There are 14 test cases, listed as following:

decode_step

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "past_state", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 1, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
past_state = np.random.randn(b, h_kv, d_k, d_v).astype(np.float32) * 0.1
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    past_state=past_state,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, past_state, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_decode_step",
    opset_imports=_OPSET,
)

delta

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "", "beta"],
    outputs=["output", "present_state"],
    update_rule="delta",
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
    update_rule="delta",
)
expect(
    node,
    inputs=[query, key, value, beta],
    outputs=[output, present_state],
    name="test_linear_attention_delta",
    opset_imports=_OPSET,
)

explicit_scale

scale = 0.25
node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=4,
    kv_num_heads=4,
    scale=scale,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
    scale=scale,
)
expect(
    node,
    inputs=[query, key, value, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_explicit_scale",
    opset_imports=_OPSET,
)

fp16

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=8,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 8, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float16)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float16), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float16)
decay = (-np.abs(np.random.randn(b, t, h_kv * d_k)) * 0.1).astype(np.float16)
beta = np.random.rand(b, t, h_kv).astype(np.float16)

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_fp16",
    opset_imports=_OPSET,
)

gated

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay"],
    outputs=["output", "present_state"],
    update_rule="gated",
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = np.random.randn(b, t, h_kv * d_k).astype(np.float32)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
# Per-key-dim decay in log-space (negative -> decay).
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
    update_rule="gated",
)
expect(
    node,
    inputs=[query, key, value, decay],
    outputs=[output, present_state],
    name="test_linear_attention_gated",
    opset_imports=_OPSET,
)

gated_delta

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_gated_delta",
    opset_imports=_OPSET,
)

gated_delta_beta_scalar

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, 1).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_gated_delta_beta_scalar",
    opset_imports=_OPSET,
)

gated_delta_gqa

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=8,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 8, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_gated_delta_gqa",
    opset_imports=_OPSET,
)

gated_delta_mqa

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=8,
    kv_num_heads=1,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 8, 1, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_gated_delta_mqa",
    opset_imports=_OPSET,
)

gated_per_head_decay

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "", "decay"],
    outputs=["output", "present_state"],
    update_rule="gated",
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = np.random.randn(b, t, h_kv * d_k).astype(np.float32)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
# Per-head scalar decay.
decay = -np.abs(np.random.randn(b, t, h_kv)).astype(np.float32) * 0.1

output, present_state = _compute(
    query,
    key,
    value,
    decay=decay,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
    update_rule="gated",
)
expect(
    node,
    inputs=[query, key, value, decay],
    outputs=[output, present_state],
    name="test_linear_attention_gated_per_head_decay",
    opset_imports=_OPSET,
)

linear

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value"],
    outputs=["output", "present_state"],
    update_rule="linear",
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = np.random.randn(b, t, h_kv * d_k).astype(np.float32)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
    update_rule="linear",
)
expect(
    node,
    inputs=[query, key, value],
    outputs=[output, present_state],
    name="test_linear_attention_linear",
    opset_imports=_OPSET,
)

linear_t1_no_past

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value"],
    outputs=["output", "present_state"],
    update_rule="linear",
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 1, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = np.random.randn(b, t, h_kv * d_k).astype(np.float32)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
    update_rule="linear",
)
expect(
    node,
    inputs=[query, key, value],
    outputs=[output, present_state],
    name="test_linear_attention_linear_t1_no_past",
    opset_imports=_OPSET,
)

no_past_explicit_zeros

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "past_state", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
past_state = np.zeros((b, h_kv, d_k, d_v), dtype=np.float32)
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    past_state=past_state,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, past_state, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_no_past_explicit_zeros",
    opset_imports=_OPSET,
)

prefill_with_past

node = onnx.helper.make_node(
    "LinearAttention",
    inputs=["query", "key", "value", "past_state", "decay", "beta"],
    outputs=["output", "present_state"],
    q_num_heads=4,
    kv_num_heads=4,
)
b, t, h_q, h_kv, d_k, d_v = 2, 4, 4, 4, 8, 8
query = np.random.randn(b, t, h_q * d_k).astype(np.float32)
key = _l2_normalize(np.random.randn(b, t, h_kv * d_k).astype(np.float32), h_kv)
value = np.random.randn(b, t, h_kv * d_v).astype(np.float32)
past_state = np.random.randn(b, h_kv, d_k, d_v).astype(np.float32) * 0.1
decay = -np.abs(np.random.randn(b, t, h_kv * d_k)).astype(np.float32) * 0.1
beta = np.random.rand(b, t, h_kv).astype(np.float32)

output, present_state = _compute(
    query,
    key,
    value,
    past_state=past_state,
    decay=decay,
    beta=beta,
    q_num_heads=h_q,
    kv_num_heads=h_kv,
)
expect(
    node,
    inputs=[query, key, value, past_state, decay, beta],
    outputs=[output, present_state],
    name="test_linear_attention_prefill_with_past",
    opset_imports=_OPSET,
)

Log

There are 1 test cases, listed as following:

log

node = onnx.helper.make_node(
    "Log",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([1, 10]).astype(np.float32)
y = np.log(x)  # expected output [0., 2.30258512]
expect(node, inputs=[x], outputs=[y], name="test_log_example")

x = np.exp(np.random.randn(3, 4, 5).astype(np.float32))
y = np.log(x)
expect(node, inputs=[x], outputs=[y], name="test_log")

LogSoftmax

There are 2 test cases, listed as following:

logsoftmax

node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
)
x = np.array([[-1, 0, 1]]).astype(np.float32)
# expected output
# [[-2.4076061 -1.407606  -0.407606 ]]
y = logsoftmax(x)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_example_1")

logsoftmax_axis

x = np.array([[0, 1, 2, 3], [10000, 10001, 10002, 10003]]).astype(np.float32)
# expected output
# [[-3.4401896  -2.4401896  -1.4401896  -0.44018966]
# [-3.4401896  -2.4401896  -1.4401896  -0.44018966]]
y = logsoftmax(x)

node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_large_number")

x = np.abs(np.random.randn(3, 4, 5).astype(np.float32))
node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
    axis=0,
)
y = logsoftmax(x, axis=0)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_axis_0")

node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
    axis=1,
)
y = logsoftmax(x, axis=1)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_axis_1")

node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
    axis=2,
)
y = logsoftmax(x, axis=2)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_axis_2")

node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
    axis=-1,
)
y = logsoftmax(x, axis=-1)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_negative_axis")

# default axis is -1
node = onnx.helper.make_node(
    "LogSoftmax",
    inputs=["x"],
    outputs=["y"],
)
expect(node, inputs=[x], outputs=[y], name="test_logsoftmax_default_axis")

Loop

There are 3 test cases, listed as following:

loop_11

# Given a tensor x of values [x1, ..., xN], and initial tensor y
# sum up its elements using a scan
# returning the final state (y+x1+x2+...+xN) as well the scan_output
# [y+x1, y+x1+x2, ..., y+x1+x2+...+xN]

y_in = onnx.helper.make_tensor_value_info("y_in", onnx.TensorProto.FLOAT, [1])
y_out = onnx.helper.make_tensor_value_info("y_out", onnx.TensorProto.FLOAT, [1])
scan_out = onnx.helper.make_tensor_value_info(
    "scan_out", onnx.TensorProto.FLOAT, [1]
)
cond_in = onnx.helper.make_tensor_value_info(
    "cond_in", onnx.TensorProto.BOOL, []
)
cond_out = onnx.helper.make_tensor_value_info(
    "cond_out", onnx.TensorProto.BOOL, []
)
iter_count = onnx.helper.make_tensor_value_info(
    "iter_count", onnx.TensorProto.INT64, []
)

x = np.array([1, 2, 3, 4, 5]).astype(np.float32)
y = np.array([-2]).astype(np.float32)

x_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["x"],
    value=onnx.helper.make_tensor(
        name="const_tensor_x",
        data_type=onnx.TensorProto.FLOAT,
        dims=x.shape,
        vals=x.flatten().astype(float),
    ),
)

one_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["one"],
    value=onnx.helper.make_tensor(
        name="const_tensor_one",
        data_type=onnx.TensorProto.INT64,
        dims=(),
        vals=[1],
    ),
)

i_add_node = onnx.helper.make_node(
    "Add", inputs=["iter_count", "one"], outputs=["end"]
)

start_unsqueeze_node = onnx.helper.make_node(
    "Unsqueeze", inputs=["iter_count"], outputs=["slice_start"], axes=[0]
)

end_unsqueeze_node = onnx.helper.make_node(
    "Unsqueeze", inputs=["end"], outputs=["slice_end"], axes=[0]
)

slice_node = onnx.helper.make_node(
    "Slice", inputs=["x", "slice_start", "slice_end"], outputs=["slice_out"]
)

y_add_node = onnx.helper.make_node(
    "Add", inputs=["y_in", "slice_out"], outputs=["y_out"]
)

identity_node = onnx.helper.make_node(
    "Identity", inputs=["cond_in"], outputs=["cond_out"]
)

scan_identity_node = onnx.helper.make_node(
    "Identity", inputs=["y_out"], outputs=["scan_out"]
)

loop_body = onnx.helper.make_graph(
    [
        identity_node,
        x_const_node,
        one_const_node,
        i_add_node,
        start_unsqueeze_node,
        end_unsqueeze_node,
        slice_node,
        y_add_node,
        scan_identity_node,
    ],
    "loop_body",
    [iter_count, cond_in, y_in],
    [cond_out, y_out, scan_out],
)

node = onnx.helper.make_node(
    "Loop",
    inputs=["trip_count", "cond", "y"],
    outputs=["res_y", "res_scan"],
    body=loop_body,
)

trip_count = np.array(5).astype(np.int64)
res_y = np.array([13]).astype(np.float32)
cond = np.array(1).astype(bool)
res_scan = np.array([-1, 1, 4, 8, 13]).astype(np.float32).reshape((5, 1))
expect(
    node,
    inputs=[trip_count, cond, y],
    outputs=[res_y, res_scan],
    name="test_loop11",
    opset_imports=[onnx.helper.make_opsetid("", 11)],
)

loop_13

# Given a tensor x of values [x1, ..., xN],
# Return a sequence of tensors of
#   [[x1], [x1, x2], ..., [x1, ..., xN]]

seq_in = onnx.helper.make_tensor_sequence_value_info(
    "seq_in", onnx.TensorProto.FLOAT, None
)
seq_out = onnx.helper.make_tensor_sequence_value_info(
    "seq_out", onnx.TensorProto.FLOAT, None
)
cond_in = onnx.helper.make_tensor_value_info(
    "cond_in", onnx.TensorProto.BOOL, []
)
cond_out = onnx.helper.make_tensor_value_info(
    "cond_out", onnx.TensorProto.BOOL, []
)
iter_count = onnx.helper.make_tensor_value_info(
    "iter_count", onnx.TensorProto.INT64, []
)

x = np.array([1, 2, 3, 4, 5]).astype(np.float32)

x_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["x"],
    value=onnx.helper.make_tensor(
        name="const_tensor_x",
        data_type=onnx.TensorProto.FLOAT,
        dims=x.shape,
        vals=x.flatten().astype(float),
    ),
)

one_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["one"],
    value=onnx.helper.make_tensor(
        name="const_tensor_one",
        data_type=onnx.TensorProto.INT64,
        dims=(),
        vals=[1],
    ),
)

zero_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["slice_start"],
    value=onnx.helper.make_tensor(
        name="const_tensor_zero",
        data_type=onnx.TensorProto.INT64,
        dims=(1,),
        vals=[0],
    ),
)

axes_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["axes"],
    value=onnx.helper.make_tensor(
        name="const_tensor_axes",
        data_type=onnx.TensorProto.INT64,
        dims=(),
        vals=[0],
    ),
)

add_node = onnx.helper.make_node(
    "Add", inputs=["iter_count", "one"], outputs=["end"]
)

end_unsqueeze_node = onnx.helper.make_node(
    "Unsqueeze", inputs=["end", "axes"], outputs=["slice_end"]
)

slice_node = onnx.helper.make_node(
    "Slice", inputs=["x", "slice_start", "slice_end"], outputs=["slice_out"]
)

insert_node = onnx.helper.make_node(
    "SequenceInsert", inputs=["seq_in", "slice_out"], outputs=["seq_out"]
)

identity_node = onnx.helper.make_node(
    "Identity", inputs=["cond_in"], outputs=["cond_out"]
)

loop_body = onnx.helper.make_graph(
    [
        identity_node,
        x_const_node,
        one_const_node,
        zero_const_node,
        add_node,
        axes_node,
        end_unsqueeze_node,
        slice_node,
        insert_node,
    ],
    "loop_body",
    [iter_count, cond_in, seq_in],
    [cond_out, seq_out],
)

node = onnx.helper.make_node(
    "Loop",
    inputs=["trip_count", "cond", "seq_empty"],
    outputs=["seq_res"],
    body=loop_body,
)

trip_count = np.array(5).astype(np.int64)
seq_empty: list[Any] = []
seq_res = [x[: int(i)] for i in x]
cond = np.array(1).astype(bool)
expect(
    node,
    inputs=[trip_count, cond, seq_empty],
    outputs=[seq_res],
    name="test_loop13_seq",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
    input_type_protos=[
        onnx.helper.make_tensor_type_proto(
            onnx.TensorProto.INT64, trip_count.shape
        ),
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.BOOL, cond.shape),
        onnx.helper.make_sequence_type_proto(
            onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, [])
        ),
    ],
)

loop_16_none

# Given a tensor sequence of values [x1, ..., xN], and an initial optional sequence of tensors [x0],
# Return a concatenated sequence of tensors of
#   [x0, [x1], [x1, x2], ..., [x1, ..., xN]]

ten_in_tp = onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, [])
seq_in_tp = onnx.helper.make_sequence_type_proto(ten_in_tp)
opt_in_tp = onnx.helper.make_optional_type_proto(seq_in_tp)
opt_in = onnx.helper.make_value_info("opt_seq_in", opt_in_tp)
seq_out = onnx.helper.make_tensor_sequence_value_info(
    "seq_out", onnx.TensorProto.FLOAT, []
)
cond_in = onnx.helper.make_tensor_value_info(
    "cond_in", onnx.TensorProto.BOOL, []
)
cond_out = onnx.helper.make_tensor_value_info(
    "cond_out", onnx.TensorProto.BOOL, []
)
iter_count = onnx.helper.make_tensor_value_info(
    "iter_count", onnx.TensorProto.INT64, []
)

x0 = np.array(0).astype(np.float32)
x = np.array([1, 2, 3, 4, 5]).astype(np.float32)

optional_has_elem_node = onnx.helper.make_node(
    "OptionalHasElement", inputs=["opt_seq_in"], outputs=["optional_has_elem"]
)

optional_is_none = onnx.helper.make_node(
    "Not", inputs=["optional_has_elem"], outputs=["optional_is_none"]
)

optional_get_elem = onnx.helper.make_node(
    "OptionalGetElement", inputs=["opt_seq_in"], outputs=["seq_in"]
)

constant_in = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["constant_in"],
    value=onnx.helper.make_tensor(
        name="const_tensor", data_type=onnx.TensorProto.FLOAT, dims=(), vals=[0]
    ),
)

seq_const_in = onnx.helper.make_node(
    "SequenceConstruct", inputs=["constant_in"], outputs=["init_seq_in"]
)

then_seq_out = onnx.helper.make_tensor_sequence_value_info(
    "init_seq_in", onnx.TensorProto.FLOAT, []
)
then_body = onnx.helper.make_graph(
    [constant_in, seq_const_in], "then_body", [], [then_seq_out]
)

else_seq_out = onnx.helper.make_tensor_sequence_value_info(
    "seq_in", onnx.TensorProto.FLOAT, []
)
else_body = onnx.helper.make_graph(
    [optional_get_elem], "else_body", [], [else_seq_out]
)

if_node = onnx.helper.make_node(
    "If",
    inputs=["optional_is_none"],
    outputs=["sequence"],
    then_branch=then_body,
    else_branch=else_body,
)

x_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["x"],
    value=onnx.helper.make_tensor(
        name="const_tensor_x",
        data_type=onnx.TensorProto.FLOAT,
        dims=x.shape,
        vals=x.flatten().astype(float),
    ),
)

one_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["one"],
    value=onnx.helper.make_tensor(
        name="const_tensor_one",
        data_type=onnx.TensorProto.INT64,
        dims=(),
        vals=[1],
    ),
)

zero_const_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["slice_start"],
    value=onnx.helper.make_tensor(
        name="const_tensor_zero",
        data_type=onnx.TensorProto.INT64,
        dims=(1,),
        vals=[0],
    ),
)

axes_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=["axes"],
    value=onnx.helper.make_tensor(
        name="const_tensor_axes",
        data_type=onnx.TensorProto.INT64,
        dims=(),
        vals=[0],
    ),
)

add_node = onnx.helper.make_node(
    "Add", inputs=["iter_count", "one"], outputs=["end"]
)

end_unsqueeze_node = onnx.helper.make_node(
    "Unsqueeze", inputs=["end", "axes"], outputs=["slice_end"]
)

slice_node = onnx.helper.make_node(
    "Slice", inputs=["x", "slice_start", "slice_end"], outputs=["slice_out"]
)

insert_node = onnx.helper.make_node(
    "SequenceInsert", inputs=["sequence", "slice_out"], outputs=["seq_out"]
)

identity_node = onnx.helper.make_node(
    "Identity", inputs=["cond_in"], outputs=["cond_out"]
)

loop_body = onnx.helper.make_graph(
    [
        identity_node,
        optional_has_elem_node,
        optional_is_none,
        if_node,
        x_const_node,
        one_const_node,
        zero_const_node,
        add_node,
        axes_node,
        end_unsqueeze_node,
        slice_node,
        insert_node,
    ],
    "loop_body",
    [iter_count, cond_in, opt_in],
    [cond_out, seq_out],
)

node = onnx.helper.make_node(
    "Loop",
    inputs=["trip_count", "cond", "opt_seq"],
    outputs=["seq_res"],
    body=loop_body,
)

trip_count = np.array(5).astype(np.int64)
cond = np.array(1).astype(bool)
seq_res = compute_loop_outputs(x, [x0], trip_count)
opt_seq_in: list[Any] = [x0]
expect(
    node,
    inputs=[trip_count, cond, opt_seq_in],
    outputs=[seq_res],
    name="test_loop16_seq_none",
    opset_imports=[onnx.helper.make_opsetid("", 16)],
    input_type_protos=[
        onnx.helper.make_tensor_type_proto(
            onnx.TensorProto.INT64, trip_count.shape
        ),
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.BOOL, cond.shape),
        opt_in_tp,
    ],
)

LpNormalization

There are 6 test cases, listed as following:

default

node = onnx.helper.make_node("LpNormalization", inputs=["x"], outputs=["y"])
x = np.array(
    [[[1.0, 2.0, 2.0], [3.0, 4.0, 0.0]], [[0.0, 5.0, 5.0], [6.0, 8.0, 0.0]]],
    dtype=np.float32,
)
lp_norm_default = np.sqrt(np.sum(x**2, axis=-1, keepdims=True))
y = x / lp_norm_default
expect(node, inputs=[x], outputs=[y], name="test_lpnormalization_default")

l1normalization_axis_0

node = onnx.helper.make_node(
    "LpNormalization", inputs=["x"], outputs=["y"], axis=0, p=1
)
x = np.array([3.0, 4.0], dtype=np.float32)
l1_norm_axis_0 = np.sum(abs(x), axis=0, keepdims=True)
y = x / l1_norm_axis_0
expect(node, inputs=[x], outputs=[y], name="test_l1normalization_axis_0")

l1normalization_axis_1

node = onnx.helper.make_node(
    "LpNormalization", inputs=["x"], outputs=["y"], axis=1, p=1
)
x = np.array([[3.0, 4.0], [6.0, 8.0]], dtype=np.float32)
l1_norm_axis_1 = np.sum(abs(x), axis=1, keepdims=True)
y = x / l1_norm_axis_1
expect(node, inputs=[x], outputs=[y], name="test_l1normalization_axis_1")

l1normalization_axis_last

node = onnx.helper.make_node(
    "LpNormalization", inputs=["x"], outputs=["y"], axis=-1, p=1
)
x = np.array(
    [[[1.0, 2.0, 2.0], [3.0, 4.0, 0.0]], [[0.0, 5.0, 5.0], [6.0, 8.0, 0.0]]],
    dtype=np.float32,
)
l1_norm_axis_last = np.sum(abs(x), axis=-1, keepdims=True)
y = x / l1_norm_axis_last
expect(node, inputs=[x], outputs=[y], name="test_l1normalization_axis_last")

l2normalization_axis_0

node = onnx.helper.make_node(
    "LpNormalization", inputs=["x"], outputs=["y"], axis=0, p=2
)
x = np.array(
    [[[1.0, 2.0, 2.0], [3.0, 4.0, 0.0]], [[0.0, 5.0, 5.0], [6.0, 8.0, 0.0]]],
    dtype=np.float32,
)
l2_norm_axis_0 = np.sqrt(np.sum(x**2, axis=0, keepdims=True))
# When norm is 0, output is 0 (0/0 = 0)
y = np.where(l2_norm_axis_0 == 0, 0, x / l2_norm_axis_0)
expect(node, inputs=[x], outputs=[y], name="test_l2normalization_axis_0")

l2normalization_axis_1

node = onnx.helper.make_node(
    "LpNormalization", inputs=["x"], outputs=["y"], axis=1, p=2
)
x = np.array([[3.0, 4.0], [6.0, 8.0]], dtype=np.float32)
l2_norm_axis_1 = np.sqrt(np.sum(x**2, axis=1, keepdims=True))
y = x / l2_norm_axis_1
expect(node, inputs=[x], outputs=[y], name="test_l2normalization_axis_1")

LpPool

There are 8 test cases, listed as following:

lppool_1d_default

"""input_shape: [1, 3, 32]
output_shape: [1, 3, 31]
"""
p = 3
kernel_shape = [2]
strides = [1]
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=kernel_shape,
    strides=strides,
    p=p,
)
x = np.random.randn(1, 3, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "LPPOOL", p=p)

expect(node, inputs=[x], outputs=[y], name="test_lppool_1d_default")

lppool_2d_default

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 31, 31]
"""
p = 4
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    p=p,
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = (2, 2)
strides = (1, 1)
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "LPPOOL", p=p)

expect(node, inputs=[x], outputs=[y], name="test_lppool_2d_default")

lppool_2d_dilations

"""input_shape: [1, 1, 4, 4]
output_shape: [1, 1, 2, 2]
"""
p = 2
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    strides=[1, 1],
    dilations=[2, 2],
    p=p,
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)

y = np.array(
    [
        [
            [
                [14.560219778561036, 16.24807680927192],
                [21.633307652783937, 23.49468024894146],
            ]
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x], outputs=[y], name="test_lppool_2d_dilations")

lppool_2d_pads

"""input_shape: [1, 3, 28, 28]
output_shape: [1, 3, 30, 30]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
p = 3
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[2, 2, 2, 2],
    p=p,
)
x = np.random.randn(1, 3, 28, 28).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (3, 3)
strides = (1, 1)
pad_bottom = pad_top = pad_right = pad_left = 2
pads = [pad_top, pad_left, pad_bottom, pad_right]
out_shape, extra_pads = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = np.pad(
    x,
    (
        (0, 0),
        (0, 0),
        (extra_pads[0], extra_pads[2]),
        (extra_pads[1], extra_pads[3]),
    ),
    mode="constant",
    constant_values=0,
)
y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "LPPOOL",
    pads_required=extra_pads,
    pads=pads,
    p=p,
)

expect(node, inputs=[x], outputs=[y], name="test_lppool_2d_pads")

lppool_2d_same_lower

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 32, 32]
pad_shape: [1, 1] -> [1, 0, 1, 0] by axis
"""
p = 4
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    auto_pad="SAME_LOWER",
    p=p,
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (2, 2)
strides = (1, 1)
out_shape = get_output_shape_auto_pad(
    "SAME_LOWER", x_shape[2:], kernel_shape, strides
)
pad_shape = get_pad_shape(
    "SAME_LOWER", x_shape[2:], kernel_shape, strides, out_shape
)
pad_bottom = pad_shape[0] // 2
pad_top = pad_shape[0] - pad_bottom
pad_right = pad_shape[1] // 2
pad_left = pad_shape[1] - pad_right
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=0,
)
pads = [pad_top, pad_left, pad_bottom, pad_right]
y = pool(
    padded, x_shape, kernel_shape, strides, out_shape, "LPPOOL", pads, pads, p=p
)

expect(node, inputs=[x], outputs=[y], name="test_lppool_2d_same_lower")

lppool_2d_same_upper

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 32, 32]
pad_shape: [1, 1] -> [0, 1, 0, 1] by axis
"""
p = 2
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    auto_pad="SAME_UPPER",
    p=p,
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (2, 2)
strides = (1, 1)
out_shape = get_output_shape_auto_pad(
    "SAME_UPPER", x_shape[2:], kernel_shape, strides
)
pad_shape = get_pad_shape(
    "SAME_UPPER", x_shape[2:], kernel_shape, strides, out_shape
)
pad_top = pad_shape[0] // 2
pad_bottom = pad_shape[0] - pad_top
pad_left = pad_shape[1] // 2
pad_right = pad_shape[1] - pad_left
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=0,
)
pads = [pad_top, pad_left, pad_bottom, pad_right]
y = pool(
    padded, x_shape, kernel_shape, strides, out_shape, "LPPOOL", pads, pads, p=p
)

expect(node, inputs=[x], outputs=[y], name="test_lppool_2d_same_upper")

lppool_2d_strides

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 10, 10]
"""
p = 2
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[5, 5],
    strides=[3, 3],
    p=p,
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = (5, 5)
strides = (3, 3)
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "LPPOOL", p=p)

expect(node, inputs=[x], outputs=[y], name="test_lppool_2d_strides")

lppool_3d_default

"""input_shape: [1, 3, 32, 32, 32]
output_shape: [1, 3, 31, 31, 31]
"""
p = 3
node = onnx.helper.make_node(
    "LpPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2, 2],
    p=p,
)
x = np.random.randn(1, 3, 32, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = [2, 2, 2]
strides = [1, 1, 1]
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "LPPOOL", p=p)

expect(node, inputs=[x], outputs=[y], name="test_lppool_3d_default")

MatMul

There are 1 test cases, listed as following:

matmul

node = onnx.helper.make_node(
    "MatMul",
    inputs=["a", "b"],
    outputs=["c"],
)

# 2d
a = np.random.randn(3, 4).astype(np.float32)
b = np.random.randn(4, 3).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_2d")

# 3d
a = np.random.randn(2, 3, 4).astype(np.float32)
b = np.random.randn(2, 4, 3).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_3d")

# 4d
a = np.random.randn(1, 2, 3, 4).astype(np.float32)
b = np.random.randn(1, 2, 4, 3).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_4d")

# broadcasting
a = np.random.randn(3, 1, 3, 4).astype(np.float32)
b = np.random.randn(1, 2, 4, 2).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_bcast")

# 1d + 3d
a = np.random.randn(4).astype(np.float32)
b = np.random.randn(2, 4, 1).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_1d_3d")

# 3d + 1d
a = np.random.randn(1, 2, 4, 3).astype(np.float32)
b = np.random.randn(3).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_4d_1d")

# 1d + 1d
a = np.random.randn(3).astype(np.float32)
b = np.random.randn(3).astype(np.float32)
c = np.matmul(a, b)
expect(node, inputs=[a, b], outputs=[c], name="test_matmul_1d_1d")

MatMulInteger

There are 1 test cases, listed as following:

matmulinteger

node = onnx.helper.make_node(
    "MatMulInteger",
    inputs=["A", "B", "a_zero_point", "b_zero_point"],
    outputs=["Y"],
)

A = np.array(
    [
        [11, 7, 3],
        [10, 6, 2],
        [9, 5, 1],
        [8, 4, 0],
    ],
    dtype=np.uint8,
)

a_zero_point = np.array([12], dtype=np.uint8)

B = np.array(
    [
        [1, 4],
        [2, 5],
        [3, 6],
    ],
    dtype=np.uint8,
)

b_zero_point = np.array([0], dtype=np.uint8)

output = np.array(
    [
        [-38, -83],
        [-44, -98],
        [-50, -113],
        [-56, -128],
    ],
    dtype=np.int32,
)

expect(
    node,
    inputs=[A, B, a_zero_point, b_zero_point],
    outputs=[output],
    name="test_matmulinteger",
)

Max

There are 2 test cases, listed as following:

max

data_0 = np.array([3, 2, 1]).astype(np.float32)
data_1 = np.array([1, 4, 4]).astype(np.float32)
data_2 = np.array([2, 5, 3]).astype(np.float32)
result = np.array([3, 5, 4]).astype(np.float32)
node = onnx.helper.make_node(
    "Max",
    inputs=["data_0", "data_1", "data_2"],
    outputs=["result"],
)
expect(
    node,
    inputs=[data_0, data_1, data_2],
    outputs=[result],
    name="test_max_example",
)

node = onnx.helper.make_node(
    "Max",
    inputs=["data_0"],
    outputs=["result"],
)
expect(node, inputs=[data_0], outputs=[data_0], name="test_max_one_input")

result = np.maximum(data_0, data_1)
node = onnx.helper.make_node(
    "Max",
    inputs=["data_0", "data_1"],
    outputs=["result"],
)
expect(
    node, inputs=[data_0, data_1], outputs=[result], name="test_max_two_inputs"
)

max_all_numeric_types

for op_dtype in all_numeric_dtypes:
    data_0 = np.array([3, 2, 1]).astype(op_dtype)
    data_1 = np.array([1, 4, 4]).astype(op_dtype)
    result = np.array([3, 4, 4]).astype(op_dtype)
    node = onnx.helper.make_node(
        "Max",
        inputs=["data_0", "data_1"],
        outputs=["result"],
    )
    expect(
        node,
        inputs=[data_0, data_1],
        outputs=[result],
        name=f"test_max_{np.dtype(op_dtype).name}",
    )

MaxPool

There are 19 test cases, listed as following:

maxpool_1d_default

"""input_shape: [1, 3, 32]
output_shape: [1, 3, 31]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2],
)
x = np.random.randn(1, 3, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = [2]
strides = [1]
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "MAX")

expect(node, inputs=[x], outputs=[y], name="test_maxpool_1d_default")

maxpool_2d_ceil

"""input_shape: [1, 1, 4, 4]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    strides=[2, 2],
    ceil_mode=True,
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[11, 12], [15, 16]]]]).astype(np.float32)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_ceil")

maxpool_2d_ceil_output_size_reduce_by_one

"""input_shape: [1, 1, 2, 2]
output_shape: [1, 1, 1, 1]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[1, 1],
    strides=[2, 2],
    ceil_mode=True,
)
x = np.array([[[[1, 2], [3, 4]]]]).astype(np.float32)
y = np.array([[[[1]]]]).astype(np.float32)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_maxpool_2d_ceil_output_size_reduce_by_one",
)

maxpool_2d_default

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 31, 31]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = (2, 2)
strides = (1, 1)
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "MAX")

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_default")

maxpool_2d_dilations

"""input_shape: [1, 1, 4, 4]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    strides=[1, 1],
    dilations=[2, 2],
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[11, 12], [15, 16]]]]).astype(np.float32)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_dilations")

maxpool_2d_pads

"""input_shape: [1, 3, 28, 28]
output_shape: [1, 3, 30, 30]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    pads=[2, 2, 2, 2],
)
x = np.random.randn(1, 3, 28, 28).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (3, 3)
strides = (1, 1)
pad_bottom = pad_top = pad_right = pad_left = 2
pads = [pad_top, pad_left, pad_bottom, pad_right]
out_shape, extra_pads = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=np.nan,
)

y = pool(
    padded,
    x_shape,
    kernel_shape,
    strides,
    out_shape,
    "MAX",
    pads_required=extra_pads,
    pads=pads,
)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_pads")

maxpool_2d_precomputed_pads

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 5, 5]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[5, 5],
    pads=[2, 2, 2, 2],
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array(
    [
        [
            [
                [13, 14, 15, 15, 15],
                [18, 19, 20, 20, 20],
                [23, 24, 25, 25, 25],
                [23, 24, 25, 25, 25],
                [23, 24, 25, 25, 25],
            ]
        ]
    ]
).astype(np.float32)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_precomputed_pads")

maxpool_2d_precomputed_same_upper

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 3, 3]
pad_shape: [2, 2] -> [1, 1, 1, 1] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[3, 3],
    strides=[2, 2],
    auto_pad="SAME_UPPER",
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[7, 9, 10], [17, 19, 20], [22, 24, 25]]]]).astype(np.float32)

expect(
    node, inputs=[x], outputs=[y], name="test_maxpool_2d_precomputed_same_upper"
)

maxpool_2d_precomputed_strides

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "MaxPool", inputs=["x"], outputs=["y"], kernel_shape=[2, 2], strides=[2, 2]
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[7, 9], [17, 19]]]]).astype(np.float32)

expect(
    node, inputs=[x], outputs=[y], name="test_maxpool_2d_precomputed_strides"
)

maxpool_2d_same_lower

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 32, 32]
pad_shape: [1, 1] -> [1, 0, 1, 0] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    auto_pad="SAME_LOWER",
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (2, 2)
strides = (1, 1)
out_shape = get_output_shape_auto_pad(
    "SAME_LOWER", x_shape[2:], kernel_shape, strides
)
pad_shape = get_pad_shape(
    "SAME_LOWER", x_shape[2:], kernel_shape, strides, out_shape
)
pad_bottom = pad_shape[0] // 2
pad_top = pad_shape[0] - pad_bottom
pad_right = pad_shape[1] // 2
pad_left = pad_shape[1] - pad_right
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=np.nan,
)
pads = [pad_top, pad_left, pad_bottom, pad_right]
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "MAX", pads, pads)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_same_lower")

maxpool_2d_same_upper

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 32, 32]
pad_shape: [1, 1] -> [0, 1, 0, 1] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2],
    auto_pad="SAME_UPPER",
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
kernel_shape = (2, 2)
strides = (1, 1)
out_shape = get_output_shape_auto_pad(
    "SAME_UPPER", x_shape[2:], kernel_shape, strides
)
pad_shape = get_pad_shape(
    "SAME_UPPER", x_shape[2:], kernel_shape, strides, out_shape
)
pad_top = pad_shape[0] // 2
pad_bottom = pad_shape[0] - pad_top
pad_left = pad_shape[1] // 2
pad_right = pad_shape[1] - pad_left
padded = np.pad(
    x,
    ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
    mode="constant",
    constant_values=np.nan,
)
pads = [pad_top, pad_left, pad_bottom, pad_right]
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "MAX", pads, pads)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_same_upper")

maxpool_2d_strides

"""input_shape: [1, 3, 32, 32]
output_shape: [1, 3, 10, 10]
"""
node = onnx.helper.make_node(
    "MaxPool", inputs=["x"], outputs=["y"], kernel_shape=[5, 5], strides=[3, 3]
)
x = np.random.randn(1, 3, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = (5, 5)
strides = (3, 3)
out_shape, pads = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "MAX")

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_strides")

maxpool_2d_uint8

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 5, 5]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[5, 5],
    pads=[2, 2, 2, 2],
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.uint8)
y = np.array(
    [
        [
            [
                [13, 14, 15, 15, 15],
                [18, 19, 20, 20, 20],
                [23, 24, 25, 25, 25],
                [23, 24, 25, 25, 25],
                [23, 24, 25, 25, 25],
            ]
        ]
    ]
).astype(np.uint8)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_2d_uint8")

maxpool_3d_default

"""input_shape: [1, 3, 32, 32, 32]
output_shape: [1, 3, 31, 31, 31]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2, 2],
)
x = np.random.randn(1, 3, 32, 32, 32).astype(np.float32)
x_shape = np.shape(x)
pads = None
kernel_shape = [2, 2, 2]
strides = [1, 1, 1]
out_shape, _ = get_output_shape_explicit_padding(
    pads, x_shape[2:], kernel_shape, strides
)
padded = x
y = pool(padded, x_shape, kernel_shape, strides, out_shape, "MAX")

expect(node, inputs=[x], outputs=[y], name="test_maxpool_3d_default")

maxpool_3d_dilations

"""input_shape: [1, 1, 4, 4, 4]
output_shape: [1, 1, 2, 2, 2]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2, 2],
    strides=[1, 1, 1],
    dilations=[2, 2, 2],
)
x = np.array(
    [
        [
            [
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[[11, 12], [15, 16]], [[11, 12], [15, 16]]]]]).astype(
    np.float32
)

expect(node, inputs=[x], outputs=[y], name="test_maxpool_3d_dilations")

maxpool_3d_dilations_use_ref_impl

"""input_shape: [1, 1, 4, 4, 4]
output_shape: [1, 1, 2, 2, 2]
"""
dilations = [2, 2, 2]
kernel_shape = [2, 2, 2]
strides = [1, 1, 1]
ceil_mode = False
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=[2, 2, 2],
    strides=[1, 1, 1],
    dilations=dilations,
)
x = np.array(
    [
        [
            [
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
                [
                    [1, 2, 3, 4],
                    [5, 6, 7, 8],
                    [9, 10, 11, 12],
                    [13, 14, 15, 16],
                ],
            ]
        ]
    ]
).astype(np.float32)

x_shape = x.shape[2:]
out_shape, pads = get_output_shape_explicit_padding(
    None, x_shape, kernel_shape, strides, dilations, ceil_mode=ceil_mode
)
padded = x
y = pool(
    padded,
    (1, 1, *x_shape),
    kernel_shape,
    strides,
    out_shape,
    "MAX",
    pads_required=pads,
    pads=None,
    dilations=dilations,
)

expect(
    node, inputs=[x], outputs=[y], name="test_maxpool_3d_dilations_use_ref_impl"
)

maxpool_3d_dilations_use_ref_impl_large

x_shape = (32, 32, 32)
dilations = (2, 2, 2)
kernel_shape = (5, 5, 5)
strides = (3, 3, 3)
ceil_mode = True

node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y"],
    kernel_shape=kernel_shape,
    strides=strides,
    dilations=dilations,
    ceil_mode=ceil_mode,
)

x = np.random.randn(1, 1, *x_shape).astype(np.float32)
out_shape, pads = get_output_shape_explicit_padding(
    None, x_shape, kernel_shape, strides, dilations, ceil_mode=ceil_mode
)
padded = np.pad(
    x,
    (
        (0, 0),
        (0, 0),
        (pads[0], pads[3]),
        (pads[1], pads[4]),
        (pads[2], pads[5]),
    ),
    mode="constant",
    constant_values=0,
)
y = pool(
    padded,
    (1, 1, *x_shape),
    kernel_shape,
    strides,
    out_shape,
    "MAX",
    pads_required=pads,
    pads=None,
    dilations=dilations,
)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_maxpool_3d_dilations_use_ref_impl_large",
)

maxpool_with_argmax_2d_precomputed_pads

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 5, 5]
pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y", "z"],
    kernel_shape=[5, 5],
    pads=[2, 2, 2, 2],
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array(
    [
        [
            [
                [13, 14, 15, 15, 15],
                [18, 19, 20, 20, 20],
                [23, 24, 25, 25, 25],
                [23, 24, 25, 25, 25],
                [23, 24, 25, 25, 25],
            ]
        ]
    ]
).astype(np.float32)
z = np.array(
    [
        [
            [
                [12, 13, 14, 14, 14],
                [17, 18, 19, 19, 19],
                [22, 23, 24, 24, 24],
                [22, 23, 24, 24, 24],
                [22, 23, 24, 24, 24],
            ]
        ]
    ]
).astype(np.int64)

expect(
    node,
    inputs=[x],
    outputs=[y, z],
    name="test_maxpool_with_argmax_2d_precomputed_pads",
)

maxpool_with_argmax_2d_precomputed_strides

"""input_shape: [1, 1, 5, 5]
output_shape: [1, 1, 2, 2]
"""
node = onnx.helper.make_node(
    "MaxPool",
    inputs=["x"],
    outputs=["y", "z"],
    kernel_shape=[2, 2],
    strides=[2, 2],
    storage_order=1,
)
x = np.array(
    [
        [
            [
                [1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25],
            ]
        ]
    ]
).astype(np.float32)
y = np.array([[[[7, 9], [17, 19]]]]).astype(np.float32)
z = np.array([[[[6, 16], [8, 18]]]]).astype(np.int64)

expect(
    node,
    inputs=[x],
    outputs=[y, z],
    name="test_maxpool_with_argmax_2d_precomputed_strides",
)

MaxUnpool

There are 2 test cases, listed as following:

with_output_shape

node = onnx.helper.make_node(
    "MaxUnpool",
    inputs=["xT", "xI", "output_shape"],
    outputs=["y"],
    kernel_shape=[2, 2],
    strides=[2, 2],
)
xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
xI = np.array([[[[5, 7], [13, 15]]]], dtype=np.int64)
output_shape = np.array((1, 1, 5, 5), dtype=np.int64)
y = np.array(
    [
        [
            [
                [0, 0, 0, 0, 0],
                [0, 5, 0, 6, 0],
                [0, 0, 0, 0, 0],
                [0, 7, 0, 8, 0],
                [0, 0, 0, 0, 0],
            ]
        ]
    ],
    dtype=np.float32,
)
expect(
    node,
    inputs=[xT, xI, output_shape],
    outputs=[y],
    name="test_maxunpool_export_with_output_shape",
)

without_output_shape

node = onnx.helper.make_node(
    "MaxUnpool",
    inputs=["xT", "xI"],
    outputs=["y"],
    kernel_shape=[2, 2],
    strides=[2, 2],
)
xT = np.array([[[[1, 2], [3, 4]]]], dtype=np.float32)
xI = np.array([[[[5, 7], [13, 15]]]], dtype=np.int64)
y = np.array(
    [[[[0, 0, 0, 0], [0, 1, 0, 2], [0, 0, 0, 0], [0, 3, 0, 4]]]],
    dtype=np.float32,
)
expect(
    node,
    inputs=[xT, xI],
    outputs=[y],
    name="test_maxunpool_export_without_output_shape",
)

Mean

There are 1 test cases, listed as following:

mean

data_0 = np.array([3, 0, 2]).astype(np.float32)
data_1 = np.array([1, 3, 4]).astype(np.float32)
data_2 = np.array([2, 6, 6]).astype(np.float32)
result = np.array([2, 3, 4]).astype(np.float32)
node = onnx.helper.make_node(
    "Mean",
    inputs=["data_0", "data_1", "data_2"],
    outputs=["result"],
)
expect(
    node,
    inputs=[data_0, data_1, data_2],
    outputs=[result],
    name="test_mean_example",
)

node = onnx.helper.make_node(
    "Mean",
    inputs=["data_0"],
    outputs=["result"],
)
expect(node, inputs=[data_0], outputs=[data_0], name="test_mean_one_input")

result = np.divide(np.add(data_0, data_1), 2.0)
node = onnx.helper.make_node(
    "Mean",
    inputs=["data_0", "data_1"],
    outputs=["result"],
)
expect(
    node, inputs=[data_0, data_1], outputs=[result], name="test_mean_two_inputs"
)

MeanVarianceNormalization

There are 1 test cases, listed as following:

meanvariancenormalization

node = onnx.helper.make_node(
    "MeanVarianceNormalization", inputs=["X"], outputs=["Y"]
)

input_data = np.array(
    [
        [
            [[0.8439683], [0.5665144], [0.05836735]],
            [[0.02916367], [0.12964272], [0.5060197]],
            [[0.79538304], [0.9411346], [0.9546573]],
        ],
        [
            [[0.17730942], [0.46192095], [0.26480448]],
            [[0.6746842], [0.01665257], [0.62473077]],
            [[0.9240844], [0.9722341], [0.11965699]],
        ],
        [
            [[0.41356155], [0.9129373], [0.59330076]],
            [[0.81929934], [0.7862604], [0.11799799]],
            [[0.69248444], [0.54119414], [0.07513223]],
        ],
    ],
    dtype=np.float32,
)

# Calculate expected output data
data_mean = np.mean(input_data, axis=(0, 2, 3), keepdims=1)
data_mean_squared = np.power(data_mean, 2)
data_squared = np.power(input_data, 2)
data_squared_mean = np.mean(data_squared, axis=(0, 2, 3), keepdims=1)
std = np.sqrt(data_squared_mean - data_mean_squared)
expected_output = (input_data - data_mean) / (std + 1e-9)

expect(node, inputs=[input_data], outputs=[expected_output], name="test_mvn")

MelWeightMatrix

There are 1 test cases, listed as following:

melweightmatrix

node = onnx.helper.make_node(
    "MelWeightMatrix",
    inputs=[
        "num_mel_bins",
        "dft_length",
        "sample_rate",
        "lower_edge_hertz",
        "upper_edge_hertz",
    ],
    outputs=["output"],
)

num_mel_bins = np.int32(8)
dft_length = np.int32(16)
sample_rate = np.int32(8192)
lower_edge_hertz = np.float32(0)
upper_edge_hertz = np.float32(8192 / 2)

num_spectrogram_bins = dft_length // 2 + 1
frequency_bins = np.arange(0, num_mel_bins + 2)

low_frequency_mel = 2595 * np.log10(1 + lower_edge_hertz / 700)
high_frequency_mel = 2595 * np.log10(1 + upper_edge_hertz / 700)
mel_step = (high_frequency_mel - low_frequency_mel) / frequency_bins.shape[0]

frequency_bins = frequency_bins * mel_step + low_frequency_mel
frequency_bins = 700 * (np.power(10, (frequency_bins / 2595)) - 1)
frequency_bins = ((dft_length + 1) * frequency_bins) // sample_rate
frequency_bins = frequency_bins.astype(int)

output = np.zeros((num_spectrogram_bins, num_mel_bins))
output.flags.writeable = True

for i in range(num_mel_bins):
    lower_frequency_value = frequency_bins[i]  # left
    center_frequency_point = frequency_bins[i + 1]  # center
    higher_frequency_point = frequency_bins[i + 2]  # right
    low_to_center = center_frequency_point - lower_frequency_value
    if low_to_center == 0:
        output[center_frequency_point, i] = 1
    else:
        for j in range(lower_frequency_value, center_frequency_point + 1):
            output[j, i] = float(j - lower_frequency_value) / float(
                low_to_center
            )
    center_to_high = higher_frequency_point - center_frequency_point
    if center_to_high > 0:
        for j in range(center_frequency_point, higher_frequency_point):
            output[j, i] = float(higher_frequency_point - j) / float(
                center_to_high
            )

# Expected output
# 1.000000, 1.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
# 0.000000, 0.000000, 1.000000, 1.000000, 0.000000, 0.000000, 0.000000, 0.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 1.000000, 0.000000, 0.000000, 0.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 1.000000, 0.000000, 0.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 1.000000, 0.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 1.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
# 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
output = output.astype(np.float32)
expect(
    node,
    inputs=[
        num_mel_bins,
        dft_length,
        sample_rate,
        lower_edge_hertz,
        upper_edge_hertz,
    ],
    outputs=[output],
    name="test_melweightmatrix",
)

Min

There are 2 test cases, listed as following:

min

data_0 = np.array([3, 2, 1]).astype(np.float32)
data_1 = np.array([1, 4, 4]).astype(np.float32)
data_2 = np.array([2, 5, 0]).astype(np.float32)
result = np.array([1, 2, 0]).astype(np.float32)
node = onnx.helper.make_node(
    "Min",
    inputs=["data_0", "data_1", "data_2"],
    outputs=["result"],
)
expect(
    node,
    inputs=[data_0, data_1, data_2],
    outputs=[result],
    name="test_min_example",
)

node = onnx.helper.make_node(
    "Min",
    inputs=["data_0"],
    outputs=["result"],
)
expect(node, inputs=[data_0], outputs=[data_0], name="test_min_one_input")

result = np.minimum(data_0, data_1)
node = onnx.helper.make_node(
    "Min",
    inputs=["data_0", "data_1"],
    outputs=["result"],
)
expect(
    node, inputs=[data_0, data_1], outputs=[result], name="test_min_two_inputs"
)

min_all_numeric_types

for op_dtype in all_numeric_dtypes:
    data_0 = np.array([3, 2, 1]).astype(op_dtype)
    data_1 = np.array([1, 4, 4]).astype(op_dtype)
    result = np.array([1, 2, 1]).astype(op_dtype)
    node = onnx.helper.make_node(
        "Min",
        inputs=["data_0", "data_1"],
        outputs=["result"],
    )
    expect(
        node,
        inputs=[data_0, data_1],
        outputs=[result],
        name=f"test_min_{np.dtype(op_dtype).name}",
    )

Mish

There are 1 test cases, listed as following:

mish

node = onnx.helper.make_node("Mish", inputs=["X"], outputs=["Y"])

input_data = np.linspace(-10, 10, 10000, dtype=np.float32)

# Calculate expected output data
expected_output = input_data * np.tanh(np.log1p(np.exp(input_data)))

expect(node, inputs=[input_data], outputs=[expected_output], name="test_mish")

Mod

There are 13 test cases, listed as following:

mod_broadcast

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.arange(0, 30).reshape([3, 2, 5]).astype(np.int32)
y = np.array([7]).astype(np.int32)
z = np.mod(x, y)
#   array([[[0, 1, 2, 3, 4],
#     [5, 6, 0, 1, 2]],

#    [[3, 4, 5, 6, 0],
#     [1, 2, 3, 4, 5]],

#    [[6, 0, 1, 2, 3],
#     [4, 5, 6, 0, 1]]], dtype=int32)
expect(node, inputs=[x, y], outputs=[z], name="test_mod_broadcast")

mod_int64_fmod

node = onnx.helper.make_node("Mod", inputs=["x", "y"], outputs=["z"], fmod=1)

x = np.array([-4, 7, 5, 4, -7, 8]).astype(np.int64)
y = np.array([2, -3, 8, -2, 3, 5]).astype(np.int64)
z = np.fmod(x, y)  # expected output [ 0,  1,  5,  0, -1,  3]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_int64_fmod")

mod_mixed_sign_float16

node = onnx.helper.make_node("Mod", inputs=["x", "y"], outputs=["z"], fmod=1)

x = np.array([-4.3, 7.2, 5.0, 4.3, -7.2, 8.0]).astype(np.float16)
y = np.array([2.1, -3.4, 8.0, -2.1, 3.4, 5.0]).astype(np.float16)
z = np.fmod(
    x, y
)  # expected output [-0.10156, 0.3984 , 5. , 0.10156, -0.3984 ,  3.]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_float16")

mod_mixed_sign_float32

node = onnx.helper.make_node("Mod", inputs=["x", "y"], outputs=["z"], fmod=1)

x = np.array([-4.3, 7.2, 5.0, 4.3, -7.2, 8.0]).astype(np.float32)
y = np.array([2.1, -3.4, 8.0, -2.1, 3.4, 5.0]).astype(np.float32)
z = np.fmod(
    x, y
)  # expected output [-0.10000038, 0.39999962, 5. , 0.10000038, -0.39999962, 3.]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_float32")

mod_mixed_sign_float64

node = onnx.helper.make_node("Mod", inputs=["x", "y"], outputs=["z"], fmod=1)

x = np.array([-4.3, 7.2, 5.0, 4.3, -7.2, 8.0]).astype(np.float64)
y = np.array([2.1, -3.4, 8.0, -2.1, 3.4, 5.0]).astype(np.float64)
z = np.fmod(x, y)  # expected output [-0.1,  0.4,  5. ,  0.1, -0.4,  3.]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_float64")

mod_mixed_sign_int16

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([-4, 7, 5, 4, -7, 8]).astype(np.int16)
y = np.array([2, -3, 8, -2, 3, 5]).astype(np.int16)
z = np.mod(x, y)  # expected output [ 0, -2,  5,  0,  2,  3]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_int16")

mod_mixed_sign_int32

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([-4, 7, 5, 4, -7, 8]).astype(np.int32)
y = np.array([2, -3, 8, -2, 3, 5]).astype(np.int32)
z = np.mod(x, y)  # expected output [ 0, -2,  5,  0,  2,  3]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_int32")

mod_mixed_sign_int64

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([-4, 7, 5, 4, -7, 8]).astype(np.int64)
y = np.array([2, -3, 8, -2, 3, 5]).astype(np.int64)
z = np.mod(x, y)  # expected output [ 0, -2,  5,  0,  2,  3]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_int64")

mod_mixed_sign_int8

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([-4, 7, 5, 4, -7, 8]).astype(np.int8)
y = np.array([2, -3, 8, -2, 3, 5]).astype(np.int8)
z = np.mod(x, y)  # expected output [ 0, -2,  5,  0,  2,  3]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_mixed_sign_int8")

mod_uint16

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([4, 7, 5]).astype(np.uint16)
y = np.array([2, 3, 8]).astype(np.uint16)
z = np.mod(x, y)  # expected output [0, 1, 5]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_uint16")

mod_uint32

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([4, 7, 5]).astype(np.uint32)
y = np.array([2, 3, 8]).astype(np.uint32)
z = np.mod(x, y)  # expected output [0, 1, 5]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_uint32")

mod_uint64

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([4, 7, 5]).astype(np.uint64)
y = np.array([2, 3, 8]).astype(np.uint64)
z = np.mod(x, y)  # expected output [0, 1, 5]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_uint64")

mod_uint8

node = onnx.helper.make_node(
    "Mod",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([4, 7, 5]).astype(np.uint8)
y = np.array([2, 3, 8]).astype(np.uint8)
z = np.mod(x, y)  # expected output [0, 1, 5]
expect(node, inputs=[x, y], outputs=[z], name="test_mod_uint8")

Momentum

There are 3 test cases, listed as following:

momentum

# Define operator attributes.
norm_coefficient = 0.001
alpha = 0.95
beta = 0.1

# Create operator.
node = onnx.helper.make_node(
    "Momentum",
    inputs=["R", "T", "X", "G", "V"],
    outputs=["X_new", "V_new"],
    norm_coefficient=norm_coefficient,
    alpha=alpha,
    beta=beta,
    mode="standard",
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar
x = np.array([1.2, 2.8], dtype=np.float32)
g = np.array([-0.94, -2.5], dtype=np.float32)
v = np.array([1.7, 3.6], dtype=np.float32)

# Compute expected outputs of Momentum.
x_new, v_new = apply_momentum(r, t, x, g, v, norm_coefficient, alpha, beta)

# Check results.
expect(
    node,
    inputs=[r, t, x, g, v],
    outputs=[x_new, v_new],
    name="test_momentum",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

momentum_multiple

# Define operator attributes.
norm_coefficient = 0.001
alpha = 0.95
beta = 0.85

node = onnx.helper.make_node(
    "Momentum",
    inputs=["R", "T", "X1", "X2", "G1", "G2", "H1", "H2"],
    outputs=["X1_new", "X2_new", "V1_new", "V2_new"],
    norm_coefficient=norm_coefficient,
    alpha=alpha,
    beta=beta,
    mode="standard",
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar

x1 = np.array([1.0], dtype=np.float32)
g1 = np.array([-1.0], dtype=np.float32)
v1 = np.array([2.0], dtype=np.float32)

x2 = np.array([1.0, 2.0], dtype=np.float32)
g2 = np.array([-1.0, -3.0], dtype=np.float32)
v2 = np.array([4.0, 1.0], dtype=np.float32)

# Compute expected outputs of Momentum.
x1_new, v1_new = apply_momentum(r, t, x1, g1, v1, norm_coefficient, alpha, beta)
x2_new, v2_new = apply_momentum(r, t, x2, g2, v2, norm_coefficient, alpha, beta)

# Check results.
expect(
    node,
    inputs=[r, t, x1, x2, g1, g2, v1, v2],
    outputs=[x1_new, x2_new, v1_new, v2_new],
    name="test_momentum_multiple",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

nesterov_momentum

# Define operator attributes.
norm_coefficient = 0.01
alpha = 0.95
beta = 1.0

# Create operator.
node = onnx.helper.make_node(
    "Momentum",
    inputs=["R", "T", "X", "G", "V"],
    outputs=["X_new", "V_new"],
    norm_coefficient=norm_coefficient,
    alpha=alpha,
    beta=beta,
    mode="nesterov",
    domain=AI_ONNX_PREVIEW_TRAINING_DOMAIN,
)

# Define operator inputs.
r = np.array(0.1, dtype=np.float32)  # scalar
t = np.array(0, dtype=np.int64)  # scalar
x = np.array([1.2, 2.8], dtype=np.float32)
g = np.array([-0.94, -2.5], dtype=np.float32)
v = np.array([1.7, 3.6], dtype=np.float32)

# Compute expected outputs of Momentum.
x_new, v_new = apply_nesterov(r, t, x, g, v, norm_coefficient, alpha, beta)

# Check results.
expect(
    node,
    inputs=[r, t, x, g, v],
    outputs=[x_new, v_new],
    name="test_nesterov_momentum",
    opset_imports=[
        onnx.helper.make_opsetid(AI_ONNX_PREVIEW_TRAINING_DOMAIN, 1)
    ],
)

Mul

There are 2 test cases, listed as following:

mul

node = onnx.helper.make_node(
    "Mul",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([4, 5, 6]).astype(np.float32)
z = x * y  # expected output [4., 10., 18.]
expect(node, inputs=[x, y], outputs=[z], name="test_mul_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul")

x = np.random.randint(4, size=(3, 4, 5), dtype=np.int8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_int8")

x = np.random.randint(4, size=(3, 4, 5), dtype=np.int16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_int16")

x = np.random.randint(4, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_uint8")

x = np.random.randint(4, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_uint16")

x = np.random.randint(4, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_uint32")

x = np.random.randint(4, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_uint64")

mul_broadcast

node = onnx.helper.make_node(
    "Mul",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
z = x * y
expect(node, inputs=[x, y], outputs=[z], name="test_mul_bcast")

Neg

There are 1 test cases, listed as following:

neg

node = onnx.helper.make_node(
    "Neg",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-4, 2]).astype(np.float32)
y = np.negative(x)  # expected output [4., -2.],
expect(node, inputs=[x], outputs=[y], name="test_neg_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.negative(x)
expect(node, inputs=[x], outputs=[y], name="test_neg")

NegativeLogLikelihoodLoss

There are 18 test cases, listed as following:

input_shape_is_NC

reduction = "none"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
)

N, C = 3, 5
np.random.seed(0)
input = np.random.rand(N, C).astype(np.float32)
target = np.random.randint(0, high=C, size=(N,)).astype(np.int64)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=None, reduction=reduction
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NC",
)

input_shape_is_NCd1

reduction = "mean"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, d1 = 3, 5, 2
np.random.seed(0)
input = np.random.rand(N, C, d1).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, d1)).astype(np.int64)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=None, reduction=reduction
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1",
)

input_shape_is_NCd1_ii

reduction = "mean"
ignore_index = np.int64(1)
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, d1 = 3, 5, 2
np.random.seed(0)
input = np.random.rand(N, C, d1).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, d1)).astype(np.int64)
target[0][0] = np.int64(1)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=None, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1_ii",
)

input_shape_is_NCd1_mean_weight_negative_ii

reduction = "mean"
ignore_index = np.int64(-1)

node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1 = 3, 5, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1)).astype(np.int64)
target[0][0] = -1
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1_mean_weight_negative_ii",
)

input_shape_is_NCd1_weight

reduction = "mean"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, d1 = 3, 5, 2
np.random.seed(0)
input = np.random.rand(N, C, d1).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, d1)).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1_weight",
)

input_shape_is_NCd1_weight_ii

reduction = "mean"
ignore_index = np.int64(1)
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, d1 = 3, 5, 2
np.random.seed(0)
input = np.random.rand(N, C, d1).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, d1)).astype(np.int64)
target[0][0] = np.int64(1)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1_weight_ii",
)

input_shape_is_NCd1d2

reduction = "none"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=None, reduction=reduction
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2",
)

input_shape_is_NCd1d2_no_weight_reduction_mean_ii

reduction = "mean"
ignore_index = np.int64(1)
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)
target[0][0][0] = np.int64(1)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_no_weight_reduction_mean_ii",
)

input_shape_is_NCd1d2_reduction_mean

reduction = "mean"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=None, reduction=reduction
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_reduction_mean",
)

input_shape_is_NCd1d2_reduction_sum

reduction = "sum"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=None, reduction=reduction
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_reduction_sum",
)

input_shape_is_NCd1d2_with_weight

reduction = "none"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_with_weight",
)

input_shape_is_NCd1d2_with_weight_reduction_mean

reduction = "mean"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_with_weight_reduction_mean",
)

input_shape_is_NCd1d2_with_weight_reduction_sum

reduction = "sum"
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_with_weight_reduction_sum",
)

input_shape_is_NCd1d2_with_weight_reduction_sum_ii

reduction = "sum"
ignore_index = np.int64(0)
node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1, dim2 = 3, 5, 6, 6
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2)).astype(np.int64)
target[0][0][0] = np.int64(0)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2_with_weight_reduction_sum_ii",
)

input_shape_is_NCd1d2d3_none_no_weight_negative_ii

reduction = "none"
ignore_index = np.int64(-5)

node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1, dim2, dim3 = 3, 5, 6, 6, 5
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2, dim3).astype(np.float32)
target = np.random.randint(0, high=C, size=(N, dim1, dim2, dim3)).astype(
    np.int64
)
target[0][0][0][0] = -5

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2d3_none_no_weight_negative_ii",
)

input_shape_is_NCd1d2d3_sum_weight_high_ii

reduction = "sum"
ignore_index = np.int64(10)

node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C = 3, 5
np.random.seed(0)
input = np.random.rand(N, C).astype(np.float32)
target = np.random.randint(0, high=C, size=(N)).astype(np.int64)
target[0] = 10
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2d3_sum_weight_high_ii",
)

input_shape_is_NCd1d2d3d4d5_mean_weight

reduction = "mean"

node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target", "weight"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2, dim3, dim4, dim5 = 3, 5, 6, 6, 5, 3, 4
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2, dim3, dim4, dim5).astype(np.float32)
target = np.random.randint(
    0, high=C, size=(N, dim1, dim2, dim3, dim4, dim5)
).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, weight=weight, reduction=reduction
)

expect(
    node,
    inputs=[input, target, weight],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2d3d4d5_mean_weight",
)

input_shape_is_NCd1d2d3d4d5_none_no_weight

reduction = "none"

node = onnx.helper.make_node(
    "NegativeLogLikelihoodLoss",
    inputs=["input", "target"],
    outputs=["loss"],
    reduction=reduction,
)

N, C, dim1, dim2, dim3, dim4, dim5 = 3, 5, 6, 6, 5, 3, 4
np.random.seed(0)
input = np.random.rand(N, C, dim1, dim2, dim3, dim4, dim5).astype(np.float32)
target = np.random.randint(
    0, high=C, size=(N, dim1, dim2, dim3, dim4, dim5)
).astype(np.int64)

negative_log_likelihood_loss = compute_negative_log_likelihood_loss(
    input, target, reduction=reduction
)

expect(
    node,
    inputs=[input, target],
    outputs=[negative_log_likelihood_loss],
    name="test_nllloss_NCd1d2d3d4d5_none_no_weight",
)

NonMaxSuppression

There are 10 test cases, listed as following:

nonmaxsuppression_center_point_box_format

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
    center_point_box=1,
)
boxes = np.array(
    [
        [
            [0.5, 0.5, 1.0, 1.0],
            [0.5, 0.6, 1.0, 1.0],
            [0.5, 0.4, 1.0, 1.0],
            [0.5, 10.5, 1.0, 1.0],
            [0.5, 10.6, 1.0, 1.0],
            [0.5, 100.5, 1.0, 1.0],
        ]
    ]
).astype(np.float32)
scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array([[0, 0, 3], [0, 0, 0], [0, 0, 5]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_center_point_box_format",
)

nonmaxsuppression_flipped_coordinates

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [1.0, 1.0, 0.0, 0.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, 0.9, 1.0, -0.1],
            [0.0, 10.0, 1.0, 11.0],
            [1.0, 10.1, 0.0, 11.1],
            [1.0, 101.0, 0.0, 100.0],
        ]
    ]
).astype(np.float32)
scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array([[0, 0, 3], [0, 0, 0], [0, 0, 5]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_flipped_coordinates",
)

nonmaxsuppression_identical_boxes

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.0, 1.0, 1.0],
        ]
    ]
).astype(np.float32)
scores = np.array(
    [[[0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9]]]
).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array([[0, 0, 0]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_identical_boxes",
)

nonmaxsuppression_iou_threshold_boundary

"""Test boundary condition where IoU exactly equals threshold.

This test verifies that the comparison is strict (>), not inclusive (>=).
When IoU exactly equals the threshold, boxes should be KEPT, not suppressed.
This follows PyTorch's NMS implementation.
"""
node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
# Two boxes with 50% overlap in each dimension
# box1=[0,0,1,1], box2=[0.5,0.5,1.5,1.5]
# Intersection area = 0.5 * 0.5 = 0.25
# Union area = 1.0 + 1.0 - 0.25 = 1.75
# IoU = 0.25 / 1.75 (exact value computed below as float32)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],  # box 0
            [0.5, 0.5, 1.5, 1.5],  # box 1 - overlaps box 0
        ]
    ]
).astype(np.float32)
scores = np.array([[[0.9, 0.8]]]).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
# Compute the exact IoU value and use it as threshold
# This ensures the threshold exactly equals the IoU
exact_iou = np.float32(0.25 / 1.75)
iou_threshold = np.array([exact_iou]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
# Both boxes should be selected because IoU == threshold (not > threshold)
selected_indices = np.array([[0, 0, 0], [0, 0, 1]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_iou_threshold_boundary",
)

nonmaxsuppression_limit_output_size

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, -0.1, 1.0, 0.9],
            [0.0, 10.0, 1.0, 11.0],
            [0.0, 10.1, 1.0, 11.1],
            [0.0, 100.0, 1.0, 101.0],
        ]
    ]
).astype(np.float32)
scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]).astype(np.float32)
max_output_boxes_per_class = np.array([2]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array([[0, 0, 3], [0, 0, 0]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_limit_output_size",
)

nonmaxsuppression_single_box

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array([[[0.0, 0.0, 1.0, 1.0]]]).astype(np.float32)
scores = np.array([[[0.9]]]).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array([[0, 0, 0]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_single_box",
)

nonmaxsuppression_suppress_by_IOU

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, -0.1, 1.0, 0.9],
            [0.0, 10.0, 1.0, 11.0],
            [0.0, 10.1, 1.0, 11.1],
            [0.0, 100.0, 1.0, 101.0],
        ]
    ]
).astype(np.float32)
scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array([[0, 0, 3], [0, 0, 0], [0, 0, 5]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_suppress_by_IOU",
)

nonmaxsuppression_suppress_by_IOU_and_scores

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, -0.1, 1.0, 0.9],
            [0.0, 10.0, 1.0, 11.0],
            [0.0, 10.1, 1.0, 11.1],
            [0.0, 100.0, 1.0, 101.0],
        ]
    ]
).astype(np.float32)
scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]).astype(np.float32)
max_output_boxes_per_class = np.array([3]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.4]).astype(np.float32)
selected_indices = np.array([[0, 0, 3], [0, 0, 0]]).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_suppress_by_IOU_and_scores",
)

nonmaxsuppression_two_batches

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, -0.1, 1.0, 0.9],
            [0.0, 10.0, 1.0, 11.0],
            [0.0, 10.1, 1.0, 11.1],
            [0.0, 100.0, 1.0, 101.0],
        ],
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, -0.1, 1.0, 0.9],
            [0.0, 10.0, 1.0, 11.0],
            [0.0, 10.1, 1.0, 11.1],
            [0.0, 100.0, 1.0, 101.0],
        ],
    ]
).astype(np.float32)
scores = np.array(
    [[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]], [[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]
).astype(np.float32)
max_output_boxes_per_class = np.array([2]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array(
    [[0, 0, 3], [0, 0, 0], [1, 0, 3], [1, 0, 0]]
).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_two_batches",
)

nonmaxsuppression_two_classes

node = onnx.helper.make_node(
    "NonMaxSuppression",
    inputs=[
        "boxes",
        "scores",
        "max_output_boxes_per_class",
        "iou_threshold",
        "score_threshold",
    ],
    outputs=["selected_indices"],
)
boxes = np.array(
    [
        [
            [0.0, 0.0, 1.0, 1.0],
            [0.0, 0.1, 1.0, 1.1],
            [0.0, -0.1, 1.0, 0.9],
            [0.0, 10.0, 1.0, 11.0],
            [0.0, 10.1, 1.0, 11.1],
            [0.0, 100.0, 1.0, 101.0],
        ]
    ]
).astype(np.float32)
scores = np.array(
    [[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3], [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]
).astype(np.float32)
max_output_boxes_per_class = np.array([2]).astype(np.int64)
iou_threshold = np.array([0.5]).astype(np.float32)
score_threshold = np.array([0.0]).astype(np.float32)
selected_indices = np.array(
    [[0, 0, 3], [0, 0, 0], [0, 1, 3], [0, 1, 0]]
).astype(np.int64)

expect(
    node,
    inputs=[
        boxes,
        scores,
        max_output_boxes_per_class,
        iou_threshold,
        score_threshold,
    ],
    outputs=[selected_indices],
    name="test_nonmaxsuppression_two_classes",
)

NonZero

There are 1 test cases, listed as following:

nonzero

node = onnx.helper.make_node(
    "NonZero",
    inputs=["condition"],
    outputs=["result"],
)

condition = np.array([[1, 0], [1, 1]], dtype=bool)
result = np.array(
    np.nonzero(condition), dtype=np.int64
)  # expected output [[0, 1, 1], [0, 0, 1]]
expect(node, inputs=[condition], outputs=[result], name="test_nonzero_example")

Not

There are 1 test cases, listed as following:

not

node = onnx.helper.make_node(
    "Not",
    inputs=["x"],
    outputs=["not"],
)

# 2d
x = (np.random.randn(3, 4) > 0).astype(bool)
expect(node, inputs=[x], outputs=[np.logical_not(x)], name="test_not_2d")

# 3d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
expect(node, inputs=[x], outputs=[np.logical_not(x)], name="test_not_3d")

# 4d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
expect(node, inputs=[x], outputs=[np.logical_not(x)], name="test_not_4d")

OneHot

There are 5 test cases, listed as following:

with_axis

axisValue = 1
on_value = 3
off_value = 1
output_type = np.float32
node = onnx.helper.make_node(
    "OneHot",
    inputs=["indices", "depth", "values"],
    outputs=["y"],
    axis=axisValue,
)
indices = np.array([[1, 9], [2, 4]], dtype=np.float32)
depth = np.float32(10)
values = np.array([off_value, on_value], dtype=output_type)
y = one_hot(indices, depth, axis=axisValue, dtype=output_type)
y = y * (on_value - off_value) + off_value
expect(
    node,
    inputs=[indices, depth, values],
    outputs=[y],
    name="test_onehot_with_axis",
)

with_negative_axis

axisValue = -2
on_value = 3
off_value = 1
output_type = np.float32
node = onnx.helper.make_node(
    "OneHot",
    inputs=["indices", "depth", "values"],
    outputs=["y"],
    axis=axisValue,
)
indices = np.array([[1, 9], [2, 4]], dtype=np.float32)
depth = np.float32(10)
values = np.array([off_value, on_value], dtype=output_type)
y = one_hot(indices, depth, axis=axisValue, dtype=output_type)
y = y * (on_value - off_value) + off_value
expect(
    node,
    inputs=[indices, depth, values],
    outputs=[y],
    name="test_onehot_with_negative_axis",
)

with_negative_indices

axisValue = 1
on_value = 3
off_value = 1
output_type = np.float32
node = onnx.helper.make_node(
    "OneHot",
    inputs=["indices", "depth", "values"],
    outputs=["y"],
    axis=axisValue,
)
indices = np.array([0, -7, -8], dtype=np.int64)

# print(y)
# [[3. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
#  [1. 1. 1. 3. 1. 1. 1. 1. 1. 1.]
#  [1. 1. 3. 1. 1. 1. 1. 1. 1. 1.]]

depth = np.float32(10)
values = np.array([off_value, on_value], dtype=output_type)
y = one_hot(indices, depth, axis=axisValue, dtype=output_type)
y = y * (on_value - off_value) + off_value
expect(
    node,
    inputs=[indices, depth, values],
    outputs=[y],
    name="test_onehot_negative_indices",
)

with_out_of_range_indices

axisValue = 1
on_value = 3
off_value = 1
output_type = np.float32
node = onnx.helper.make_node(
    "OneHot",
    inputs=["indices", "depth", "values"],
    outputs=["y"],
    axis=axisValue,
)
# Indices outside [-depth, depth-1] map to an all-off_value row.
indices = np.array([5, -6, -1], dtype=np.int64)

# print(y)
# [[1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 3.]]

depth = np.float32(5)
values = np.array([off_value, on_value], dtype=output_type)
y = one_hot(indices, depth, axis=axisValue, dtype=output_type)
y = y * (on_value - off_value) + off_value
expect(
    node,
    inputs=[indices, depth, values],
    outputs=[y],
    name="test_onehot_out_of_range_indices",
)

without_axis

on_value = 5
off_value = 2
output_type = np.int32
node = onnx.helper.make_node(
    "OneHot", inputs=["indices", "depth", "values"], outputs=["y"]
)
indices = np.array([0, 7, 8], dtype=np.int64)
depth = np.float32(12)
values = np.array([off_value, on_value], dtype=output_type)
y = one_hot(indices, depth, dtype=output_type)
y = y * (on_value - off_value) + off_value
expect(
    node,
    inputs=[indices, depth, values],
    outputs=[y],
    name="test_onehot_without_axis",
)

OptionalHasElement

There are 4 test cases, listed as following:

empty

optional = None

tensor_type_proto = onnx.helper.make_tensor_type_proto(
    elem_type=onnx.TensorProto.INT32, shape=[]
)
optional_type_proto = onnx.helper.make_optional_type_proto(tensor_type_proto)

# OptionalHasElement takes a tensor or optional as input
for input_type_proto in [tensor_type_proto, optional_type_proto]:
    input_name_options = {
        "empty": "optional_input",
        "empty_no_input_name": "",
        "empty_no_input": None,
    }
    for test_name_surfix, input_name in input_name_options.items():
        if input_type_proto == tensor_type_proto and input_name:
            # the input tensor cannot be empty if input name is provided.
            continue
        node = onnx.helper.make_node(
            "OptionalHasElement",
            inputs=[] if input_name is None else [input_name],
            outputs=["output"],
        )
        output = optional_has_element_reference_implementation(optional)
        test_name = (
            "test_optional_has_element_"
            + test_name_surfix
            + (
                "_optional_input"
                if input_type_proto == optional_type_proto
                else "_tensor_input"
            )
        )
        expect(
            node,
            inputs=[optional] if input_name else [],
            outputs=[output],
            input_type_protos=[input_type_proto] if input_name else [],
            name=test_name,
        )

get_element_sequence

optional = [np.array([1, 2, 3, 4]).astype(np.int32)]
tensor_type_proto = onnx.helper.make_tensor_type_proto(
    elem_type=onnx.TensorProto.INT32,
    shape=[
        4,
    ],
)
seq_type_proto = onnx.helper.make_sequence_type_proto(tensor_type_proto)
optional_type_proto = onnx.helper.make_optional_type_proto(seq_type_proto)

node = onnx.helper.make_node(
    "OptionalGetElement", inputs=["optional_input"], outputs=["output"]
)
output = optional_get_element_reference_implementation(optional)
expect(
    node,
    inputs=[optional],
    outputs=[output],
    input_type_protos=[optional_type_proto],
    name="test_optional_get_element_optional_sequence",
)
expect(
    node,
    inputs=[optional],
    outputs=[output],
    input_type_protos=[seq_type_proto],
    name="test_optional_get_element_sequence",
)

get_element_tensor

optional = np.array([1, 2, 3, 4]).astype(np.float32)
tensor_type_proto = onnx.helper.make_tensor_type_proto(
    elem_type=onnx.TensorProto.FLOAT,
    shape=[
        4,
    ],
)
optional_type_proto = onnx.helper.make_optional_type_proto(tensor_type_proto)

node = onnx.helper.make_node(
    "OptionalGetElement", inputs=["optional_input"], outputs=["output"]
)
output = optional_get_element_reference_implementation(optional)
expect(
    node,
    inputs=[optional],
    outputs=[output],
    input_type_protos=[optional_type_proto],
    name="test_optional_get_element_optional_tensor",
)
expect(
    node,
    inputs=[optional],
    outputs=[output],
    input_type_protos=[tensor_type_proto],
    name="test_optional_get_element_tensor",
)

optionalhaselement

optional = np.array([1, 2, 3, 4]).astype(np.float32)
tensor_type_proto = onnx.helper.make_tensor_type_proto(
    elem_type=onnx.TensorProto.FLOAT,
    shape=[
        4,
    ],
)
optional_type_proto = onnx.helper.make_optional_type_proto(tensor_type_proto)

# OptionalHasElement takes a tensor or optional as input
for input_type_protos in [tensor_type_proto, optional_type_proto]:
    node = onnx.helper.make_node(
        "OptionalHasElement", inputs=["optional_input"], outputs=["output"]
    )
    output = optional_has_element_reference_implementation(optional)
    test_name = "test_optional_has_element_" + (
        "optional_input"
        if input_type_protos == optional_type_proto
        else "tensor_input"
    )
    expect(
        node,
        inputs=[optional],
        outputs=[output],
        input_type_protos=[optional_type_proto],
        name=test_name,
    )

Or

There are 2 test cases, listed as following:

or

node = onnx.helper.make_node(
    "Or",
    inputs=["x", "y"],
    outputs=["or"],
)

# 2d
x = (np.random.randn(3, 4) > 0).astype(bool)
y = (np.random.randn(3, 4) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or2d")

# 3d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(3, 4, 5) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or3d")

# 4d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or4d")

or_broadcast

node = onnx.helper.make_node(
    "Or",
    inputs=["x", "y"],
    outputs=["or"],
)

# 3d vs 1d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(5) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or_bcast3v1d")

# 3d vs 2d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(4, 5) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or_bcast3v2d")

# 4d vs 2d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(5, 6) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or_bcast4v2d")

# 4d vs 3d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(4, 5, 6) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or_bcast4v3d")

# 4d vs 4d
x = (np.random.randn(1, 4, 1, 6) > 0).astype(bool)
y = (np.random.randn(3, 1, 5, 6) > 0).astype(bool)
z = np.logical_or(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_or_bcast4v4d")

PRelu

There are 2 test cases, listed as following:

prelu

node = onnx.helper.make_node(
    "PRelu",
    inputs=["x", "slope"],
    outputs=["y"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
slope = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, 0, np.inf) + np.clip(x, -np.inf, 0) * slope

expect(node, inputs=[x, slope], outputs=[y], name="test_prelu_example")

prelu_broadcast

node = onnx.helper.make_node(
    "PRelu",
    inputs=["x", "slope"],
    outputs=["y"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
slope = np.random.randn(5).astype(np.float32)
y = np.clip(x, 0, np.inf) + np.clip(x, -np.inf, 0) * slope

expect(node, inputs=[x, slope], outputs=[y], name="test_prelu_broadcast")

Pad

There are 4 test cases, listed as following:

constant_pad

node = onnx.helper.make_node(
    "Pad", inputs=["x", "pads", "value"], outputs=["y"], mode="constant"
)
x = np.random.randn(1, 3, 4, 5).astype(np.float32)
pads = np.array([0, 0, 1, 3, 0, 0, 2, 4]).astype(
    np.int64
)  # pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
value = np.float32(1.2)
y = pad_impl(x, pads, "constant", 1.2)

expect(node, inputs=[x, pads, value], outputs=[y], name="test_constant_pad")

constant_pad_axes

node = onnx.helper.make_node(
    "Pad", inputs=["x", "pads", "value", "axes"], outputs=["y"], mode="constant"
)
x = np.random.randn(1, 3, 4, 5).astype(np.float32)
pads = np.array([0, 3, 0, 4]).astype(
    np.int64
)  # pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
value = np.float32(1.2)
axes = np.array([1, 3], dtype=np.int64)
y = pad_impl(
    x,
    pads,
    "constant",
    1.2,
    [1, 3],
)

expect(
    node,
    inputs=[x, pads, value, axes],
    outputs=[y],
    name="test_constant_pad_axes",
)

constant_pad_negative_axes

node = onnx.helper.make_node(
    "Pad", inputs=["x", "pads", "value", "axes"], outputs=["y"], mode="constant"
)
x = np.random.randn(1, 3, 4, 5).astype(np.float32)
pads = np.array([0, 3, 0, 4]).astype(
    np.int64
)  # pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
value = np.float32(1.2)
axes = np.array([-3, -1], dtype=np.int64)
y = pad_impl(
    x,
    pads,
    "constant",
    1.2,
    [-3, -1],
)

expect(
    node,
    inputs=[x, pads, value, axes],
    outputs=[y],
    name="test_constant_pad_negative_axes",
)

reflection_edge_and_wrap_pad

for mode in ("edge", "reflect", "wrap"):
    node = onnx.helper.make_node(
        "Pad", inputs=["x", "pads"], outputs=["y"], mode=mode
    )
    x = np.random.randn(1, 3, 4, 5).astype(np.int32)
    pads = np.array([0, 0, 1, 1, 0, 0, 1, 1]).astype(
        np.int64
    )  # pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
    y = pad_impl(x, pads, mode)

    expect(node, inputs=[x, pads], outputs=[y], name=f"test_{mode}_pad")

Pow

There are 3 test cases, listed as following:

pow

node = onnx.helper.make_node(
    "Pow",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([4, 5, 6]).astype(np.float32)
z = pow(x, y)  # expected output [1., 32., 729.]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_example")

x = np.arange(60).reshape(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = pow(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_pow")

pow_broadcast

node = onnx.helper.make_node(
    "Pow",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array(2).astype(np.float32)
z = pow(x, y)  # expected output [1., 4., 9.]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_bcast_scalar")

node = onnx.helper.make_node(
    "Pow",
    inputs=["x", "y"],
    outputs=["z"],
)
x = np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32)
y = np.array([1, 2, 3]).astype(np.float32)
# expected output [[1, 4, 27], [4, 25, 216]]
z = pow(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_pow_bcast_array")

types

node = onnx.helper.make_node(
    "Pow",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([4, 5, 6]).astype(np.int64)
z = pow(x, y)  # expected output [1., 32., 729.]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_float32_int64")

x = np.array([1, 2, 3]).astype(np.int64)
y = np.array([4, 5, 6]).astype(np.float32)
z = pow(x, y)  # expected output [1, 32, 729]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_int64_float32")

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([4, 5, 6]).astype(np.int32)
z = pow(x, y)  # expected output [1., 32., 729.]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_float32_int32")

x = np.array([1, 2, 3]).astype(np.int32)
y = np.array([4, 5, 6]).astype(np.float32)
z = pow(x, y)  # expected output [1, 32, 729]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_int32_float32")

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([4, 5, 6]).astype(np.uint64)
z = pow(x, y)  # expected output [1., 32., 729.]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_float32_uint64")

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([4, 5, 6]).astype(np.uint32)
z = pow(x, y)  # expected output [1., 32., 729.]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_float32_uint32")

x = np.array([1, 2, 3]).astype(np.int64)
y = np.array([4, 5, 6]).astype(np.int64)
z = pow(x, y)  # expected output [1, 32, 729]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_int64_int64")

x = np.array([1, 2, 3]).astype(np.int32)
y = np.array([4, 5, 6]).astype(np.int32)
z = pow(x, y)  # expected output [1, 32, 729]
expect(node, inputs=[x, y], outputs=[z], name="test_pow_types_int32_int32")

QLinearConv

There are 1 test cases, listed as following:

qlinearconv

node = onnx.helper.make_node(
    "QLinearConv",
    inputs=[
        "x",
        "x_scale",
        "x_zero_point",
        "w",
        "w_scale",
        "w_zero_point",
        "y_scale",
        "y_zero_point",
    ],
    outputs=["y"],
)

x = np.array(
    [
        [255, 174, 162, 25, 203, 168, 58],
        [15, 59, 237, 95, 129, 0, 64],
        [56, 242, 153, 221, 168, 12, 166],
        [232, 178, 186, 195, 237, 162, 237],
        [188, 39, 124, 77, 80, 102, 43],
        [127, 230, 21, 83, 41, 40, 134],
        [255, 154, 92, 141, 42, 148, 247],
    ],
    dtype=np.uint8,
).reshape((1, 1, 7, 7))

x_scale = np.float32(0.00369204697)
x_zero_point = np.uint8(132)

w = np.array([0], dtype=np.uint8).reshape((1, 1, 1, 1))

w_scale = np.array([0.00172794575], dtype=np.float32)
w_zero_point = np.array([255], dtype=np.uint8)

y_scale = np.float32(0.00162681262)
y_zero_point = np.uint8(123)

output = np.array(
    [
        [0, 81, 93, 230, 52, 87, 197],
        [240, 196, 18, 160, 126, 255, 191],
        [199, 13, 102, 34, 87, 243, 89],
        [23, 77, 69, 60, 18, 93, 18],
        [67, 216, 131, 178, 175, 153, 212],
        [128, 25, 234, 172, 214, 215, 121],
        [0, 101, 163, 114, 213, 107, 8],
    ],
    dtype=np.uint8,
).reshape((1, 1, 7, 7))

expect(
    node,
    inputs=[
        x,
        x_scale,
        x_zero_point,
        w,
        w_scale,
        w_zero_point,
        y_scale,
        y_zero_point,
    ],
    outputs=[output],
    name="test_qlinearconv",
)

QLinearMatMul

There are 1 test cases, listed as following:

int

for quant_type_name in ["uint8", "int8"]:
    quant_type = getattr(np, quant_type_name)
    for dtype_name in ["float32", "float16"]:
        dtype = getattr(np, dtype_name)
        node = onnx.helper.make_node(
            "QLinearMatMul",
            inputs=[
                "a",
                "a_scale",
                "a_zero_point",
                "b",
                "b_scale",
                "b_zero_point",
                "y_scale",
                "y_zero_point",
            ],
            outputs=["y"],
        )

        # 2D
        a = np.array([[208, 236, 0, 238], [3, 214, 255, 29]])
        if quant_type == np.int8:
            a -= 127
        a = a.astype(quant_type)

        a_scale = np.array([0.0066], dtype=dtype)
        a_zero_point = np.array(
            [113 - 127] if quant_type == np.int8 else [113], dtype=quant_type
        )

        b = np.array(
            [[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]]
        )
        if quant_type == np.int8:
            b -= 127
        b = b.astype(quant_type)

        b_scale = np.array([0.00705], dtype=dtype)
        b_zero_point = np.array(
            [114 - 127] if quant_type == np.int8 else [114], dtype=quant_type
        )

        y_scale = np.array([0.0107], dtype=dtype)
        y_zero_point = np.array(
            [118 - 127] if quant_type == np.int8 else [118], dtype=quant_type
        )

        if quant_type == np.int8:
            output = np.array([[41, -12, -9], [1, -75, -128]])
        else:
            output = np.array([[168, 115, 255], [1, 66, 151]])
        output = output.astype(quant_type)

        expect(
            node,
            inputs=[
                a,
                a_scale,
                a_zero_point,
                b,
                b_scale,
                b_zero_point,
                y_scale,
                y_zero_point,
            ],
            outputs=[output],
            name=f"test_qlinearmatmul_2D_{quant_type_name}_{dtype_name}",
        )

        # 3D
        a = np.array(
            [
                [[208, 236, 0, 238], [3, 214, 255, 29]],
                [[208, 236, 0, 238], [3, 214, 255, 29]],
            ],
        )
        if quant_type == np.int8:
            a -= 127
        a = a.astype(quant_type)

        a_scale = np.array([0.0066], dtype=dtype)
        a_zero_point = np.array(
            [113 - 127] if quant_type == np.int8 else [113], dtype=quant_type
        )

        b = np.array(
            [
                [[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]],
                [[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]],
            ],
        )
        if quant_type == np.int8:
            b -= 127
        b = b.astype(quant_type)

        b_scale = np.array([0.00705], dtype=dtype)
        b_zero_point = np.array([114], dtype=quant_type)

        y_scale = np.array([0.0107], dtype=dtype)
        y_zero_point = np.array(
            [118 - 127] if quant_type == np.int8 else [118], dtype=quant_type
        )

        if quant_type == np.int8:
            output = np.array(
                [
                    [[-86, -128, -128], [115, 39, -121]],
                    [[-86, -128, -128], [115, 39, -121]],
                ]
            )
        else:
            output = np.array(
                [
                    [[168, 115, 255], [1, 66, 151]],
                    [[168, 115, 255], [1, 66, 151]],
                ]
            )
        output = output.astype(quant_type)

        expect(
            node,
            inputs=[
                a,
                a_scale,
                a_zero_point,
                b,
                b_scale,
                b_zero_point,
                y_scale,
                y_zero_point,
            ],
            outputs=[output],
            name=f"test_qlinearmatmul_3D_{quant_type_name}_{dtype_name}",
        )

QuantizeLinear

There are 13 test cases, listed as following:

axis

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
)

x = np.array(
    [
        [
            [[-162, 10], [-100, 232], [-20, -50]],
            [[-76, 0], [0, 252], [32, -44]],
            [[245, -485], [-960, -270], [-375, -470]],
        ],
    ],
    dtype=np.float32,
)
y_scale = np.array([2, 4, 5], dtype=np.float32)
y_zero_point = np.array([84, 24, 196], dtype=np.uint8)
y = (x / y_scale.reshape(1, 3, 1, 1) + y_zero_point.reshape(1, 3, 1, 1)).astype(
    np.uint8
)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_axis",
)

blocked_asymmetric

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
    axis=1,
    block_size=2,
)

x = np.array(
    [
        [6.0, 12.0, 50.0, 5.0],
        [1.0, 8.0, 4.0, 5.0],
        [0.0, 20.0, 10.0, 4.0],
    ],
    dtype=np.float32,
)
y_scale = np.array(
    [
        [1.5, 2.5],
        [3.0, 4.9],
        [5.1, 6.9],
    ],
    dtype=np.float32,
)
y_zero_point = np.array(
    [
        [0, 1],
        [1, 0],
        [2, 3],
    ],
    dtype=np.uint8,
)
# x.shape = (3, 4)
# y_scale.shape = (3, 2)
assert y_scale.shape == y_zero_point.shape
block_axis = 1
# The block shape is [x.shape[i] // y_scale.shape[i] for i in range(len(x.shape))] = (1, 2)
assert all(
    x.shape[i] == y_scale.shape[i]
    for i in range(len(x.shape))
    if i != block_axis
)
assert x.shape[block_axis] % y_scale.shape[block_axis] == 0
repeats = x.shape[block_axis] // y_scale.shape[block_axis]

# Create element-wise scale and zero point
y_scale_elementwise = np.repeat(y_scale, repeats=repeats, axis=block_axis)
y_zero_point_elementwise = np.repeat(
    y_zero_point, repeats=repeats, axis=block_axis
)

y = np.rint(x / y_scale_elementwise + y_zero_point_elementwise).astype(np.uint8)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_blocked_asymmetric",
)

blocked_symmetric

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale"],
    outputs=["y"],
    axis=1,
    block_size=2,
    output_dtype=TensorProto.INT16,
)

x = np.array(
    [
        [6.0, -8, -10, 5.0],
        [1.0, 8.0, 4.0, 5.0],
        [0.0, 20.0, 10.0, 4.0],
    ],
    dtype=np.float32,
)

y_scale = np.array(
    [
        [1.5, 2.5],
        [3.0, 4.9],
        [5.1, 6.9],
    ],
    dtype=np.float32,
)

# x.shape = (3, 4)
# y_scale.shape = (3, 2)

block_axis = 1
# The block shape is [x.shape[i] // y_scale.shape[i] for i in range(len(x.shape))] = (1, 2)
assert all(
    x.shape[i] == y_scale.shape[i]
    for i in range(len(x.shape))
    if i != block_axis
)
assert x.shape[block_axis] % y_scale.shape[block_axis] == 0
repeats = x.shape[block_axis] // y_scale.shape[block_axis]

# Create element-wise scale and zero point
y_scale_elementwise = np.repeat(y_scale, repeats=repeats, axis=block_axis)

y_val = np.clip(
    np.rint(x / y_scale_elementwise), a_min=-32768, a_max=32767
).astype(np.int16)
y = make_tensor(
    "y",
    TensorProto.INT16,
    x.shape,
    y_val,
)
expect(
    node,
    inputs=[x, y_scale],
    outputs=[y],
    name="test_quantizelinear_blocked_symmetric",
)

e4m3fn

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
)

x = np.array([0.0, 1.0, 2.0, 100000.0, 200.0]).astype(np.float32)
y_scale = np.float32(2)
y_zero_point = make_tensor("y_zero_point", TensorProto.FLOAT8E4M3FN, [1], [0])
y = make_tensor("y", TensorProto.FLOAT8E4M3FN, [5], [0, 0.5, 1, 448, 96])

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_e4m3fn",
)

e5m2

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
)

x = np.array([0.0, 1.0, 2.0, 100000.0, 200.0]).astype(np.float32)
y_scale = np.float32(2)
y_zero_point = make_tensor("y_zero_point", TensorProto.FLOAT8E5M2, [1], [0.0])
y = make_tensor("y", TensorProto.FLOAT8E5M2, [5], [0, 0.5, 1, 49152, 96])

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_e5m2",
)

float4e2m1

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
    axis=0,
)

x = np.array(
    [
        [0.0, 2.5, 4.8, 8.6],
        [-30, -20, 6, 9],
        [-0.0, -2.5, -4.8, -8.6],
    ]
).astype(np.float32)

y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
y_zero_point = make_tensor(
    "y_zero_point",
    TensorProto.FLOAT4E2M1,
    y_scale.shape,
    np.zeros_like(y_scale),
)
y = make_tensor(
    "y",
    TensorProto.FLOAT4E2M1,
    x.shape,
    [0, 1, 2, 4, -6, -6, 2, 3, 0, -0.5, -1, -2],
)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_float4e2m1",
)

int16

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
)

x = np.array(
    [
        0.0,
        -514.0,
        3.0,
        -3.0,
        2.9,
        -2.9,
        3.1,
        -3.1,
        65022.0,
        -66046.0,
        65023.0,
        -66047.0,
        65024.0,
        -66048.0,
        70000.0,
        -70000.0,
    ]
).astype(np.float32)
y_scale = np.float32(2.0)
y_zero_point = np.int16(256)
y = np.array(
    [
        256,
        -1,
        258,
        254,
        257,
        255,
        258,
        254,
        32767,
        -32767,
        32767,
        -32768,
        32767,
        -32768,
        32767,
        -32768,
    ]
).astype(np.int16)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_int16",
)

int2

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
    axis=0,
)
x = np.array(
    [
        [0.0, 2.5, 4.8, 8.6],
        [-4.0, -3.0, 1.0, 2.0],
        [-0.0, -2.5, -4.8, -8.6],
    ],
    dtype=np.float32,
)
y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
y_zero_point = make_tensor(
    "y_zero_point", TensorProto.INT2, y_scale.shape, np.zeros_like(y_scale)
)
y = make_tensor(
    "y", TensorProto.INT2, x.shape, [0, 1, 1, 1, -1, -1, 0, 1, 0, -1, -1, -2]
)
expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_int2",
)

int4

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
    axis=0,
)

x = np.array(
    [
        [0.0, 2.5, 4.8, 8.6],
        [-30, -20, 6, 9],
        [12, 15, 16, 40],
    ]
).astype(np.float32)

y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
y_zero_point = make_tensor(
    "y_zero_point", TensorProto.INT4, y_scale.shape, np.ones_like(y_scale)
)
y = make_tensor(
    "y", TensorProto.INT4, x.shape, [1, 2, 3, 5, -8, -6, 3, 4, 4, 5, 5, 7]
)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_int4",
)

quantizelinear

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
)

x = np.array([0, 2, 3, 1000, -254, -1000]).astype(np.float32)
y_scale = np.float32(2)
y_zero_point = np.uint8(128)
y = np.array([128, 129, 130, 255, 1, 0]).astype(np.uint8)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear",
)

uint16

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
)

x = np.array(
    [
        0.0,
        -128.0,
        3.0,
        -3.0,
        2.9,
        -2.9,
        3.1,
        -3.1,
        65536.0,
        -65534.0,
        70000.0,
        -70000.0,
    ]
).astype(np.float32)
y_scale = np.float32(2.0)
y_zero_point = np.uint16(32767)
y = np.array(
    [
        32767,
        32703,
        32769,
        32765,
        32768,
        32766,
        32769,
        32765,
        65535,
        0,
        65535,
        0,
    ]
).astype(np.uint16)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_uint16",
)

uint2

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
    axis=0,
)

x = np.array(
    [
        [0.0, 2.5, 4.8, 8.6],
        [-2.0, -1.0, 1.0, 3.0],
        [4.0, 5.0, 6.0, 7.0],
    ],
    dtype=np.float32,
)
y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
y_zero_point = make_tensor(
    "y_zero_point", TensorProto.UINT2, y_scale.shape, np.zeros_like(y_scale)
)
y = make_tensor(
    "y", TensorProto.UINT2, x.shape, [0, 1, 2, 3, 0, 0, 0, 1, 1, 1, 2, 2]
)
expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_uint2",
)

uint4

node = onnx.helper.make_node(
    "QuantizeLinear",
    inputs=["x", "y_scale", "y_zero_point"],
    outputs=["y"],
    axis=0,
)

x = np.array(
    [
        [0.0, 2.5, 4.8, 8.6],
        [-30, -20, 6, 9],
        [12, 15, 16, 40],
    ]
).astype(np.float32)

y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
y_zero_point = make_tensor(
    "y_zero_point", TensorProto.UINT4, y_scale.shape, np.ones_like(y_scale)
)
y = make_tensor(
    "y", TensorProto.UINT4, x.shape, [1, 2, 3, 5, 0, 0, 3, 4, 4, 5, 5, 11]
)

expect(
    node,
    inputs=[x, y_scale, y_zero_point],
    outputs=[y],
    name="test_quantizelinear_uint4",
)

RMSNormalization

There are 4 test cases, listed as following:

d

X = np.random.randn(3, 4).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    Y = _rms_normalization(X, W, axis=axis)

    node = onnx.helper.make_node(
        "RMSNormalization",
        inputs=["X", "W"],
        outputs=["Y"],
        axis=axis,
    )

    if axis < 0:
        name = f"test_rms_normalization_2d_axis_negative_{-axis}"
    else:
        name = f"test_rms_normalization_2d_axis{axis}"

    expect(node, inputs=[X, W], outputs=[Y], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

d_epsilon

epsilon = 1e-1
X = np.random.randn(2, 3, 5).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    Y = _rms_normalization(X, W, axis=axis, epsilon=epsilon)
    node = onnx.helper.make_node(
        "RMSNormalization",
        inputs=["X", "W"],
        outputs=["Y"],
        axis=axis,
        epsilon=epsilon,
    )

    if axis < 0:
        name = f"test_rms_normalization_3d_axis_negative_{-axis}_epsilon"
    else:
        name = f"test_rms_normalization_3d_axis{axis}_epsilon"

    expect(node, inputs=[X, W], outputs=[Y], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

default_axis

X = np.random.randn(2, 3, 4, 5).astype(np.float32)

# Default axis in RMSNormalization is -1.
normalized_shape = calculate_normalized_shape(X.shape, -1)
W = np.random.randn(*normalized_shape).astype(np.float32)
# Axis is default to -1 in the reference implementation.
Y = _rms_normalization(X, W)

# Not specifying axis attribute means -1.
node = onnx.helper.make_node(
    "RMSNormalization",
    inputs=["X", "W"],
    outputs=["Y"],
)

expect(
    node,
    inputs=[X, W],
    outputs=[Y],
    name="test_rms_normalization_default_axis",
)

rmsnormalization

X = np.random.randn(2, 3, 4, 5).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    Y = _rms_normalization(X, W, axis=axis)

    node = onnx.helper.make_node(
        "RMSNormalization",
        inputs=["X", "W"],
        outputs=["Y"],
        axis=axis,
    )

    if axis < 0:
        name = f"test_rms_normalization_4d_axis_negative_{-axis}"
    else:
        name = f"test_rms_normalization_4d_axis{axis}"

    expect(node, inputs=[X, W], outputs=[Y], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

RNN

There are 4 test cases, listed as following:

batchwise

input = np.array([[[1.0, 2.0]], [[3.0, 4.0]], [[5.0, 6.0]]]).astype(np.float32)

input_size = 2
hidden_size = 4
weight_scale = 0.5
layout = 1

node = onnx.helper.make_node(
    "RNN",
    inputs=["X", "W", "R"],
    outputs=["Y", "Y_h"],
    hidden_size=hidden_size,
    layout=layout,
)

W = weight_scale * np.ones((1, hidden_size, input_size)).astype(np.float32)
R = weight_scale * np.ones((1, hidden_size, hidden_size)).astype(np.float32)

rnn = RNNHelper(X=input, W=W, R=R, layout=layout)
Y, Y_h = rnn.step()
expect(
    node,
    inputs=[input, W, R],
    outputs=[Y.astype(np.float32), Y_h.astype(np.float32)],
    name="test_simple_rnn_batchwise",
)

defaults

input = np.array([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]).astype(np.float32)

input_size = 2
hidden_size = 4
weight_scale = 0.1

node = onnx.helper.make_node(
    "RNN", inputs=["X", "W", "R"], outputs=["", "Y_h"], hidden_size=hidden_size
)

W = weight_scale * np.ones((1, hidden_size, input_size)).astype(np.float32)
R = weight_scale * np.ones((1, hidden_size, hidden_size)).astype(np.float32)

rnn = RNNHelper(X=input, W=W, R=R)
_, Y_h = rnn.step()
expect(
    node,
    inputs=[input, W, R],
    outputs=[Y_h.astype(np.float32)],
    name="test_simple_rnn_defaults",
)

initial_bias

input = np.array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]).astype(
    np.float32
)

input_size = 3
hidden_size = 5
custom_bias = 0.1
weight_scale = 0.1

node = onnx.helper.make_node(
    "RNN",
    inputs=["X", "W", "R", "B"],
    outputs=["", "Y_h"],
    hidden_size=hidden_size,
)

W = weight_scale * np.ones((1, hidden_size, input_size)).astype(np.float32)
R = weight_scale * np.ones((1, hidden_size, hidden_size)).astype(np.float32)

# Adding custom bias
W_B = custom_bias * np.ones((1, hidden_size)).astype(np.float32)
R_B = np.zeros((1, hidden_size)).astype(np.float32)
B = np.concatenate((W_B, R_B), axis=1)

rnn = RNNHelper(X=input, W=W, R=R, B=B)
_, Y_h = rnn.step()
expect(
    node,
    inputs=[input, W, R, B],
    outputs=[Y_h.astype(np.float32)],
    name="test_simple_rnn_with_initial_bias",
)

seq_length

input = np.array(
    [
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
        [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]],
    ]
).astype(np.float32)

input_size = 3
hidden_size = 5

node = onnx.helper.make_node(
    "RNN",
    inputs=["X", "W", "R", "B"],
    outputs=["", "Y_h"],
    hidden_size=hidden_size,
)

W = np.random.randn(1, hidden_size, input_size).astype(np.float32)
R = np.random.randn(1, hidden_size, hidden_size).astype(np.float32)

# Adding custom bias
W_B = np.random.randn(1, hidden_size).astype(np.float32)
R_B = np.random.randn(1, hidden_size).astype(np.float32)
B = np.concatenate((W_B, R_B), axis=1)

rnn = RNNHelper(X=input, W=W, R=R, B=B)
_, Y_h = rnn.step()
expect(
    node,
    inputs=[input, W, R, B],
    outputs=[Y_h.astype(np.float32)],
    name="test_rnn_seq_length",
)

Range

There are 4 test cases, listed as following:

range_bfloat16_type_positive_delta

node = onnx.helper.make_node(
    "Range",
    inputs=["start", "limit", "delta"],
    outputs=["output"],
)

start = np.array(1.0, dtype=ml_dtypes.bfloat16)
limit = np.array(5.0, dtype=ml_dtypes.bfloat16)
delta = np.array(2.0, dtype=ml_dtypes.bfloat16)

output = np.arange(1.0, 5.0, 2.0, dtype=np.float32).astype(
    ml_dtypes.bfloat16
)  # expected output [1.0, 3.0] as bfloat16

expect(
    node,
    inputs=[start, limit, delta],
    outputs=[output],
    name="test_range_bfloat16_type_positive_delta",
)

range_float16_type_positive_delta

node = onnx.helper.make_node(
    "Range",
    inputs=["start", "limit", "delta"],
    outputs=["output"],
)

start = np.float16(1)
limit = np.float16(5)
delta = np.float16(2)

output = np.arange(
    start, limit, delta, dtype=np.float16
)  # expected output [1.0, 3.0]
expect(
    node,
    inputs=[start, limit, delta],
    outputs=[output],
    name="test_range_float16_type_positive_delta",
)

range_float_type_positive_delta

node = onnx.helper.make_node(
    "Range",
    inputs=["start", "limit", "delta"],
    outputs=["output"],
)

start = np.float32(1)
limit = np.float32(5)
delta = np.float32(2)

output = np.arange(
    start, limit, delta, dtype=np.float32
)  # expected output [1.0, 3.0]
expect(
    node,
    inputs=[start, limit, delta],
    outputs=[output],
    name="test_range_float_type_positive_delta",
)

range_int32_type_negative_delta

node = onnx.helper.make_node(
    "Range",
    inputs=["start", "limit", "delta"],
    outputs=["output"],
)

start = np.int32(10)
limit = np.int32(6)
delta = np.int32(-3)

output = np.arange(
    start, limit, delta, dtype=np.int32
)  # expected output [10, 7]
expect(
    node,
    inputs=[start, limit, delta],
    outputs=[output],
    name="test_range_int32_type_negative_delta",
)

Reciprocal

There are 1 test cases, listed as following:

reciprocal

node = onnx.helper.make_node(
    "Reciprocal",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-4, 2]).astype(np.float32)
y = np.reciprocal(x)  # expected output [-0.25, 0.5],
expect(node, inputs=[x], outputs=[y], name="test_reciprocal_example")

x = np.random.rand(3, 4, 5).astype(np.float32) + 0.5
y = np.reciprocal(x)
expect(node, inputs=[x], outputs=[y], name="test_reciprocal")

ReduceL1

There are 5 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = np.array([], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceL1", inputs=["data", "axes"], outputs=["reduced"], keepdims=keepdims
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sum(a=np.abs(data), axis=None, keepdims=keepdims == 1)
# print(reduced)
# [[[78.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(a=np.abs(data), axis=None, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([2], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceL1",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sum(a=np.abs(data), axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[3., 7.], [11., 15.], [19., 23.]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(a=np.abs(data), axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_do_not_keepdims_random",
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceL1",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
reduced = np.array(np.zeros(reduced_shape, dtype=np.float32))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceL1",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sum(a=np.abs(data), axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[3.], [7.]], [[11.], [15.]], [[19.], [23.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_keep_dims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(a=np.abs(data), axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_keep_dims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceL1",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sum(a=np.abs(data), axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[3.], [7.]], [[11.], [15.]], [[19.], [23.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_negative_axes_keep_dims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(a=np.abs(data), axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l1_negative_axes_keep_dims_random",
)

ReduceL2

There are 5 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = np.array([], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceL2", inputs=["data", "axes"], outputs=["reduced"], keepdims=keepdims
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sqrt(np.sum(a=np.square(data), axis=None, keepdims=keepdims == 1))
# print(reduced)
# [[[25.49509757]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sqrt(np.sum(a=np.square(data), axis=None, keepdims=keepdims == 1))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([2], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceL2",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sqrt(
    np.sum(a=np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
)
# print(reduced)
# [[2.23606798, 5.],
# [7.81024968, 10.63014581],
# [13.45362405, 16.2788206]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sqrt(
    np.sum(a=np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_do_not_keepdims_random",
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceL2",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
reduced = np.array(np.zeros(reduced_shape, dtype=np.float32))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceL2",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sqrt(
    np.sum(a=np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
)
# print(reduced)
# [[[2.23606798], [5.]]
# [[7.81024968], [10.63014581]]
# [[13.45362405], [16.2788206 ]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_keep_dims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sqrt(
    np.sum(a=np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_keep_dims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceL2",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.reshape(np.arange(1, np.prod(shape) + 1, dtype=np.float32), shape)
# print(data)
# [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]], [[9., 10.], [11., 12.]]]

reduced = np.sqrt(
    np.sum(a=np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
)
# print(reduced)
# [[[2.23606798], [5.]]
# [[7.81024968], [10.63014581]]
# [[13.45362405], [16.2788206 ]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_negative_axes_keep_dims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sqrt(
    np.sum(a=np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_l2_negative_axes_keep_dims_random",
)

ReduceLogSum

There are 4 test cases, listed as following:

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceLogSum",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
zero = np.array(np.zeros(reduced_shape, dtype=np.float32))
reduced = np.log(zero)  # -inf

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_empty_set",
)

keepdims

node = onnx.helper.make_node(
    "ReduceLogSum", inputs=["data", "axes"], outputs=["reduced"]
)
data = np.random.ranf([3, 4, 5]).astype(np.float32)
reduced = np.log(np.sum(data, keepdims=True))
axes = np.array([], dtype=np.int64)
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_default",
)

negative_axes_keepdims

axes = np.array([-2], dtype=np.int64)
node = onnx.helper.make_node(
    "ReduceLogSum", inputs=["data", "axes"], outputs=["reduced"]
)
data = np.random.ranf([3, 4, 5]).astype(np.float32)
reduced = np.log(np.sum(data, axis=tuple(axes), keepdims=True))
# print(reduced)
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_negative_axes",
)

nokeepdims

shape = [3, 4, 5]
axes = np.array([2, 1], dtype=np.int64)

node = onnx.helper.make_node(
    "ReduceLogSum",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=0,
)
data = np.random.ranf(shape).astype(np.float32)
reduced = np.log(np.sum(data, axis=tuple(axes), keepdims=False))
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_desc_axes",
)

axes = np.array([0, 1], dtype=np.int64)
node = onnx.helper.make_node(
    "ReduceLogSum",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=0,
)
data = np.random.ranf(shape).astype(np.float32)
reduced = np.log(np.sum(data, axis=tuple(axes), keepdims=False))
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_asc_axes",
)

ReduceLogSumExp

There are 5 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = np.array([], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceLogSumExp",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]], dtype=np.double
)
reduced = np.log(np.sum(np.exp(data), axis=None, keepdims=keepdims == 1))
# print(reduced)
# [[[60.00671387]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.double)
reduced = np.log(np.sum(np.exp(data), axis=None, keepdims=keepdims == 1))
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0
node = onnx.helper.make_node(
    "ReduceLogSumExp",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]], dtype=np.double
)
reduced = np.log(np.sum(np.exp(data), axis=tuple(axes), keepdims=keepdims == 1))
# print(reduced)
# [[20., 2.31326175]
# [40.00004578, 2.31326175]
# [60.00671387, 2.31326175]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.double)
reduced = np.log(np.sum(np.exp(data), axis=tuple(axes), keepdims=keepdims == 1))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_do_not_keepdims_random",
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceLogSumExp",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
zero = np.array(np.zeros(reduced_shape, dtype=np.float32))
reduced = np.log(zero)  # -inf

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1
node = onnx.helper.make_node(
    "ReduceLogSumExp",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]], dtype=np.double
)
reduced = np.log(np.sum(np.exp(data), axis=tuple(axes), keepdims=keepdims == 1))
# print(reduced)
# [[[20., 2.31326175]]
# [[40.00004578, 2.31326175]]
# [[60.00671387, 2.31326175]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.double)
reduced = np.log(np.sum(np.exp(data), axis=tuple(axes), keepdims=keepdims == 1))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_keepdims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1
node = onnx.helper.make_node(
    "ReduceLogSumExp",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]], dtype=np.double
)
reduced = np.log(np.sum(np.exp(data), axis=tuple(axes), keepdims=keepdims == 1))
# print(reduced)
# [[[20., 2.31326175]]
# [[40.00004578, 2.31326175]]
# [[60.00671387, 2.31326175]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_negative_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.double)
reduced = np.log(
    np.sum(np.exp(data), axis=tuple(axes.tolist()), keepdims=keepdims == 1)
)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_log_sum_exp_negative_axes_keepdims_random",
)

ReduceMax

There are 6 test cases, listed as following:

bool_inputs

axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMax",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[True, True], [True, False], [False, True], [False, False]],
)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=bool(keepdims))
# print(reduced)
# [[True],
#  [True],
#  [True],
#  [False]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_bool_inputs",
)

default_axes_keepdims

shape = [3, 2, 2]
axes = None
keepdims = 1
node = onnx.helper.make_node(
    "ReduceMax", inputs=["data"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.maximum.reduce(data, axis=axes, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data],
    outputs=[reduced],
    name="test_reduce_max_default_axes_keepdim_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.maximum.reduce(data, axis=axes, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data],
    outputs=[reduced],
    name="test_reduce_max_default_axes_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceMax",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[20., 2.]
# [40., 2.]
# [60., 2.]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_do_not_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_do_not_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceMax",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
one = np.array(np.ones(reduced_shape, dtype=np.float32))
zero = np.array(np.zeros(reduced_shape, dtype=np.float32))
reduced = -(one / zero)  # -inf

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMax",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[20., 2.]]
# [[40., 2.]]
# [[60., 2.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMax",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[20., 2.]]
# [[40., 2.]]
# [[60., 2.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_negative_axes_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.maximum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_max_negative_axes_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

ReduceMean

There are 4 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = np.array([], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMean",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.mean(data, axis=None, keepdims=keepdims == 1)
# print(reduced)
# [[[18.25]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.mean(data, axis=None, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceMean",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.mean(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[12.5, 1.5]
# [35., 1.5]
# [57.5, 1.5]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.mean(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_do_not_keepdims_random",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMean",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.mean(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[12.5, 1.5]]
# [[35., 1.5]]
# [[57.5, 1.5]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.mean(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_keepdims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMean",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.mean(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[12.5, 1.5]]
# [[35., 1.5]]
# [[57.5, 1.5]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_negative_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.mean(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_mean_negative_axes_keepdims_random",
)

ReduceMin

There are 6 test cases, listed as following:

bool_inputs

axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMin",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[True, True], [True, False], [False, True], [False, False]],
)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=bool(keepdims))
# print(reduced)
# [[ True],
#  [False],
#  [False],
#  [False]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_bool_inputs",
)

default_axes_keepdims

shape = [3, 2, 2]
axes = None
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMin", inputs=["data"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.minimum.reduce(data, axis=axes, keepdims=keepdims == 1)
# print(reduced)
# [[[1.]]]

expect(
    node,
    inputs=[data],
    outputs=[reduced],
    name="test_reduce_min_default_axes_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.minimum.reduce(data, axis=axes, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data],
    outputs=[reduced],
    name="test_reduce_min_default_axes_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceMin",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[5., 1.]
# [30., 1.]
# [55., 1.]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_do_not_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_do_not_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceMin",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
one = np.array(np.ones(reduced_shape, dtype=np.float32))
zero = np.array(np.zeros(reduced_shape, dtype=np.float32))
reduced = one / zero  # inf

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMin",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[5., 1.]]
# [[30., 1.]]
# [[55., 1.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceMin",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]],
    dtype=np.float32,
)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[5., 1.]]
# [[30., 1.]]
# [[55., 1.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_negative_axes_keepdims_example",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.minimum.reduce(data, axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_min_negative_axes_keepdims_random",
    opset_imports=[onnx.helper.make_opsetid("", 18)],
)

ReduceProd

There are 5 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = None
keepdims = 1

node = onnx.helper.make_node(
    "ReduceProd", inputs=["data"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.prod(data, axis=axes, keepdims=keepdims == 1)
# print(reduced)
# [[[4.790016e+08]]]

expect(
    node,
    inputs=[data],
    outputs=[reduced],
    name="test_reduce_prod_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.prod(data, axis=axes, keepdims=keepdims == 1)
expect(
    node,
    inputs=[data],
    outputs=[reduced],
    name="test_reduce_prod_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceProd",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.prod(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[3., 8.]
# [35., 48.]
# [99., 120.]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.prod(data, axis=tuple(axes), keepdims=keepdims == 1)
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_do_not_keepdims_random",
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceProd",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
reduced = np.array(np.ones(reduced_shape, dtype=np.float32))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceProd",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.prod(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[3., 8.]]
# [[35., 48.]]
# [[99., 120.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.prod(data, axis=tuple(axes), keepdims=keepdims == 1)
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_keepdims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceProd",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.prod(data, axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[3., 8.]]
# [[35., 48.]]
# [[99., 120.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_negative_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.prod(data, axis=tuple(axes), keepdims=keepdims == 1)
expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_prod_negative_axes_keepdims_random",
)

ReduceSum

There are 7 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = np.array([], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSum", inputs=["data", "axes"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(data, axis=None, keepdims=keepdims == 1)
# print(reduced)
# [[[78.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(data, axis=None, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceSum", inputs=["data", "axes"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(data, axis=tuple(axes.tolist()), keepdims=keepdims == 1)
# print(reduced)
# [[4., 6.]
# [12., 14.]
# [20., 22.]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(data, axis=tuple(axes.tolist()), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_do_not_keepdims_random",
)

empty_axes_input_noop

shape = [3, 2, 2]
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSum",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
    noop_with_empty_axes=True,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
axes = np.array([], dtype=np.int64)
reduced = np.array(data)
# print(reduced)
# [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_empty_axes_input_noop_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.array(data)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_empty_axes_input_noop",
)

empty_set

"""Test case with the reduced-axis of size zero."""
shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceSum",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
reduced = np.array(np.zeros(reduced_shape, dtype=np.float32))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSum", inputs=["data", "axes"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(data, axis=tuple(axes.tolist()), keepdims=keepdims == 1)
# print(reduced)
# [[[4., 6.]]
# [[12., 14.]]
# [[20., 22.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(data, axis=tuple(axes.tolist()), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_keepdims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSum", inputs=["data", "axes"], outputs=["reduced"], keepdims=keepdims
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(data, axis=tuple(axes.tolist()), keepdims=keepdims == 1)
# print(reduced)
# [[[4., 6.]]
# [[12., 14.]]
# [[20., 22.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_negative_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(data, axis=tuple(axes.tolist()), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_negative_axes_keepdims_random",
)

non_reduced_axis_zero

"""Test case with the non-reduced-axis of size zero."""
shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 0, 1]

node = onnx.helper.make_node(
    "ReduceSum",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([2], dtype=np.int64)
reduced = np.array([], dtype=np.float32).reshape(reduced_shape)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_empty_set_non_reduced_axis_zero",
)

ReduceSumSquare

There are 5 test cases, listed as following:

default_axes_keepdims

shape = [3, 2, 2]
axes = np.array([], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSumSquare",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(np.square(data), axis=None, keepdims=keepdims == 1)
# print(reduced)
# [[[650.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_default_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(np.square(data), axis=None, keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_default_axes_keepdims_random",
)

do_not_keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 0

node = onnx.helper.make_node(
    "ReduceSumSquare",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[10., 20.]
# [74., 100.]
# [202., 244.]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_do_not_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(np.square(data), axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_do_not_keepdims_random",
)

empty_set

shape = [2, 0, 4]
keepdims = 1
reduced_shape = [2, 1, 4]

node = onnx.helper.make_node(
    "ReduceSumSquare",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array([], dtype=np.float32).reshape(shape)
axes = np.array([1], dtype=np.int64)
reduced = np.array(np.zeros(reduced_shape, dtype=np.float32))

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_empty_set",
)

keepdims

shape = [3, 2, 2]
axes = np.array([1], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSumSquare",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[10., 20.]]
# [[74., 100.]]
# [[202., 244.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(np.square(data), axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_keepdims_random",
)

negative_axes_keepdims

shape = [3, 2, 2]
axes = np.array([-2], dtype=np.int64)
keepdims = 1

node = onnx.helper.make_node(
    "ReduceSumSquare",
    inputs=["data", "axes"],
    outputs=["reduced"],
    keepdims=keepdims,
)

data = np.array(
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=np.float32
)
reduced = np.sum(np.square(data), axis=tuple(axes), keepdims=keepdims == 1)
# print(reduced)
# [[[10., 20.s]]
# [[74., 100.]]
# [[202., 244.]]]

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_negative_axes_keepdims_example",
)

np.random.seed(0)
data = np.random.uniform(-10, 10, shape).astype(np.float32)
reduced = np.sum(np.square(data), axis=tuple(axes), keepdims=keepdims == 1)

expect(
    node,
    inputs=[data, axes],
    outputs=[reduced],
    name="test_reduce_sum_square_negative_axes_keepdims_random",
)

RegexFullMatch

There are 3 test cases, listed as following:

basic

node = onnx.helper.make_node(
    "RegexFullMatch",
    inputs=["X"],
    outputs=["Y"],
    pattern=r"www\.[\w.-]+\.\bcom\b",
)

x = np.array(["www.google.com", "www.facebook.com", "www.bbc.co.uk"]).astype(
    object
)
result = np.array([True, True, False])
expect(node, inputs=[x], outputs=[result], name="test_regex_full_match_basic")

match_email_domain

node = onnx.helper.make_node(
    "RegexFullMatch",
    inputs=["X"],
    outputs=["Y"],
    pattern=r"(\W|^)[\w.\-]{0,25}@(yahoo|gmail)\.com(\W|$)",
)

x = np.array(
    [
        ["account@gmail.com", "account@hotmail.com"],
        ["not email", "account2@yahoo.com"],
    ]
).astype(object)
result = np.array([[True, False], [False, True]])
expect(
    node,
    inputs=[x],
    outputs=[result],
    name="test_regex_full_match_email_domain",
)

match_empty

node = onnx.helper.make_node(
    "RegexFullMatch",
    inputs=["X"],
    outputs=["Y"],
    pattern=r"(\W|^)[\w.\-]{0,25}@(yahoo|gmail)\.com(\W|$)",
)

x = np.array([[], []]).astype(object)
result = np.array([[], []]).astype(bool)
expect(
    node,
    inputs=[x],
    outputs=[result],
    name="test_regex_full_match_empty",
)

Relu

There are 1 test cases, listed as following:

relu

node = onnx.helper.make_node(
    "Relu",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, 0, np.inf)

expect(node, inputs=[x], outputs=[y], name="test_relu")

Reshape

There are 2 test cases, listed as following:

allowzero

original_shape = [0, 3, 4]
test_cases = {
    "allowzero_reordered": np.array([3, 4, 0], dtype=np.int64),
}
data = np.random.random_sample(original_shape).astype(np.float32)

for test_name, shape in test_cases.items():
    node = onnx.helper.make_node(
        "Reshape",
        inputs=["data", "shape"],
        outputs=["reshaped"],
        allowzero=1,  # if allowzero=1, final shape = (3, 4, 0)
        # if allowzero=0, final shape = (3, 4, 4)
    )

    reshaped = reshape_reference_implementation(data, shape, allowzero=1)

    expect(
        node,
        inputs=[data, shape],
        outputs=[reshaped],
        name="test_reshape_" + test_name,
    )

reshape

original_shape = [2, 3, 4]
test_cases = {
    "reordered_all_dims": np.array([4, 2, 3], dtype=np.int64),
    "reordered_last_dims": np.array([2, 4, 3], dtype=np.int64),
    "reduced_dims": np.array([2, 12], dtype=np.int64),
    "extended_dims": np.array([2, 3, 2, 2], dtype=np.int64),
    "one_dim": np.array([24], dtype=np.int64),
    "negative_dim": np.array([2, -1, 2], dtype=np.int64),
    "negative_extended_dims": np.array([-1, 2, 3, 4], dtype=np.int64),
    "zero_dim": np.array([2, 0, 4, 1], dtype=np.int64),
    "zero_and_negative_dim": np.array([2, 0, 1, -1], dtype=np.int64),
}
data = np.random.random_sample(original_shape).astype(np.float32)

for test_name, shape in test_cases.items():
    node = onnx.helper.make_node(
        "Reshape",
        inputs=["data", "shape"],
        outputs=["reshaped"],
    )

    reshaped = reshape_reference_implementation(data, shape)

    expect(
        node,
        inputs=[data, shape],
        outputs=[reshaped],
        name="test_reshape_" + test_name,
    )

Resize

There are 39 test cases, listed as following:

resize_downsample_scales_cubic

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.8, 0.8], dtype=np.float32)

# [[[[ 1.47119141  2.78125     4.08251953]
#    [ 6.71142578  8.02148438  9.32275391]
#    [11.91650391 13.2265625  14.52783203]]]]
output = interpolate_nd(
    data, lambda x, _: cubic_coeffs(x), scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_cubic",
)

resize_downsample_scales_cubic_A_n0p5_exclude_outside

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
    cubic_coeff_a=-0.5,
    exclude_outside=True,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.8, 0.8], dtype=np.float32)

# [[[[ 1.36812675  2.6695014   4.0133367 ]
#    [ 6.57362535  7.875       9.2188353 ]
#    [11.94896657 13.25034122 14.59417652]]]]
output = interpolate_nd(
    data,
    lambda x, _: cubic_coeffs(x, A=-0.5),
    scale_factors=scales,
    exclude_outside=True,
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_cubic_A_n0p5_exclude_outside",
)

resize_downsample_scales_cubic_align_corners

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
    coordinate_transformation_mode="align_corners",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.8, 0.8], dtype=np.float32)

# [[[[ 1.          2.39519159  3.79038317]
#    [ 6.58076634  7.97595793  9.37114951]
#    [12.16153268 13.55672427 14.95191585]]]]
output = interpolate_nd(
    data,
    lambda x, _: cubic_coeffs(x),
    scale_factors=scales,
    coordinate_transformation_mode="align_corners",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_cubic_align_corners",
)

resize_downsample_scales_cubic_antialias

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
    antialias=1,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)

# [[[[ 2.5180721  4.2858863]
#    [ 9.589329  11.357142 ]]]]
output = interpolate_nd(
    data, cubic_coeffs_antialias, scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_cubic_antialias",
)

resize_downsample_scales_linear

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)

# [[[[2.6666665 4.3333331]]]]
output = interpolate_nd(
    data, lambda x, _: linear_coeffs(x), scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_linear",
)

resize_downsample_scales_linear_align_corners

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="align_corners",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)

# [[[[1.       3.142857]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    scale_factors=scales,
    coordinate_transformation_mode="align_corners",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_linear_align_corners",
)

resize_downsample_scales_linear_antialias

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
    antialias=1,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)

# [[[[ 2.875  4.5  ]
#    [ 9.375 11.   ]]]]
output = interpolate_nd(
    data, linear_coeffs_antialias, scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_linear_antialias",
)

resize_downsample_scales_linear_half_pixel_symmetric

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="half_pixel_symmetric",
)

data = np.array([[[[1, 2, 3, 4]]]], dtype=np.float32)
scales = np.array([1.0, 1.0, 1.0, 0.6], dtype=np.float32)

# [[[[1.6666667, 3.3333333]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    scale_factors=scales,
    coordinate_transformation_mode="half_pixel_symmetric",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_linear_half_pixel_symmetric",
)

resize_downsample_scales_nearest

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="nearest",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)

# [[[[1. 3.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_downsample_scales_nearest",
)

resize_downsample_sizes_cubic

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="cubic",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 3, 3], dtype=np.int64)

# [[[[ 1.63078704  3.00462963  4.37847222]
#    [ 7.12615741  8.5         9.87384259]
#    [12.62152778 13.99537037 15.36921296]]]]
output = interpolate_nd(
    data, lambda x, _: cubic_coeffs(x), output_size=sizes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_cubic",
)

resize_downsample_sizes_cubic_antialias

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="cubic",
    antialias=1,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 3, 3], dtype=np.int64)

# [[[[ 1.7750092  3.1200073  4.4650054]
#    [ 7.1550016  8.5        9.844998 ]
#    [12.534994  13.8799925 15.224991 ]]]]
output = interpolate_nd(data, cubic_coeffs_antialias, output_size=sizes).astype(
    np.float32
)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_cubic_antialias",
)

resize_downsample_sizes_linear_antialias

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="linear",
    antialias=1,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 3, 3], dtype=np.int64)

# [[[[ 2.3636363  3.590909   4.818182 ]
#    [ 7.2727275  8.5        9.727273 ]
#    [12.181818  13.409091  14.636364 ]]]]
output = interpolate_nd(
    data, linear_coeffs_antialias, output_size=sizes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_linear_antialias",
)

resize_downsample_sizes_linear_pytorch_half_pixel

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="pytorch_half_pixel",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 3, 1], dtype=np.int64)

# [[[[ 1.6666666]
#    [ 7.       ]
#    [12.333333 ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    output_size=sizes,
    coordinate_transformation_mode="pytorch_half_pixel",
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_linear_pytorch_half_pixel",
)

resize_downsample_sizes_nearest

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 1, 3], dtype=np.int64)

# [[[[1. 2. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), output_size=sizes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_nearest",
)

resize_downsample_sizes_nearest_not_larger

keep_aspect_ratio_policy = "not_larger"
axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 3], dtype=np.int64)  # Results in 1x2

# [[[[1. 3.]]]]
output = interpolate_nd(
    data,
    lambda x, _: nearest_coeffs(x),
    output_size=sizes,
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_nearest_not_larger",
)

resize_downsample_sizes_nearest_not_smaller

keep_aspect_ratio_policy = "not_smaller"
axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 3], dtype=np.int64)  # Results in 2x3

# [[[[1. 2. 4.]
#    [5. 6. 8.]]]]
output = interpolate_nd(
    data,
    lambda x, _: nearest_coeffs(x),
    output_size=sizes,
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_downsample_sizes_nearest_not_smaller",
)

resize_tf_crop_and_resize

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "roi", "", "sizes"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="tf_crop_and_resize",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

# Note: for some rois, the result may be different with that of TF for inaccurate floating point
roi = np.array([0, 0, 0.4, 0.6, 1, 1, 0.6, 0.8], dtype=np.float32)
sizes = np.array([1, 1, 3, 3], dtype=np.int64)

# [[[[ 7.6000004  7.9        8.2      ]
#    [ 8.8        9.1        9.400001 ]
#    [10.        10.3       10.6      ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    output_size=sizes,
    roi=roi,
    coordinate_transformation_mode="tf_crop_and_resize",
).astype(np.float32)

expect(
    node,
    inputs=[data, roi, sizes],
    outputs=[output],
    name="test_resize_tf_crop_and_resize",
)

resize_tf_crop_and_resize_axes_2_3

axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "roi", "", "sizes"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="tf_crop_and_resize",
    axes=axes,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

# Note: for some rois, the result may be different with that of TF for inaccurate floating point
roi = np.array([0.4, 0.6, 0.6, 0.8], dtype=np.float32)
sizes = np.array([3, 3], dtype=np.int64)

# [[[[ 7.6000004  7.9        8.2      ]
#    [ 8.8        9.1        9.400001 ]
#    [10.        10.3       10.6      ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    output_size=sizes,
    roi=roi,
    axes=axes,
    coordinate_transformation_mode="tf_crop_and_resize",
).astype(np.float32)

expect(
    node,
    inputs=[data, roi, sizes],
    outputs=[output],
    name="test_resize_tf_crop_and_resize_axes_2_3",
)

resize_tf_crop_and_resize_axes_3_2

axes = [3, 2]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "roi", "", "sizes"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="tf_crop_and_resize",
    axes=axes,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

# Note: for some rois, the result may be different with that of TF for inaccurate floating point
roi = np.array([0.6, 0.4, 0.8, 0.6], dtype=np.float32)
sizes = np.array([3, 3], dtype=np.int64)

# [[[[ 7.6000004  7.9        8.2      ]
#    [ 8.8        9.1        9.400001 ]
#    [10.        10.3       10.6      ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    output_size=sizes,
    roi=roi,
    axes=axes,
    coordinate_transformation_mode="tf_crop_and_resize",
).astype(np.float32)

expect(
    node,
    inputs=[data, roi, sizes],
    outputs=[output],
    name="test_resize_tf_crop_and_resize_axes_3_2",
)

resize_tf_crop_and_resize_extrapolation_value

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "roi", "", "sizes"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="tf_crop_and_resize",
    extrapolation_value=10.0,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

# Note: for some rois, the result may be different with that of TF for inaccurate floating point
roi = np.array([0, 0, 0.4, 0.6, 1, 1, 1.2, 1.7], dtype=np.float32)
sizes = np.array([1, 1, 3, 3], dtype=np.int64)

# [[[[ 7.6000004 10.        10.       ]
#    [12.400001  10.        10.       ]
#    [10.        10.        10.       ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    output_size=sizes,
    roi=roi,
    coordinate_transformation_mode="tf_crop_and_resize",
    extrapolation_value=10.0,
).astype(np.float32)

expect(
    node,
    inputs=[data, roi, sizes],
    outputs=[output],
    name="test_resize_tf_crop_and_resize_extrapolation_value",
)

resize_upsample_scales_cubic

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)

# [[[[ 0.47265625  0.76953125  1.24609375  1.875       2.28125
#      2.91015625  3.38671875  3.68359375]
#    [ 1.66015625  1.95703125  2.43359375  3.0625      3.46875
#      4.09765625  4.57421875  4.87109375]
#    [ 3.56640625  3.86328125  4.33984375  4.96875     5.375
#      6.00390625  6.48046875  6.77734375]
#    [ 6.08203125  6.37890625  6.85546875  7.484375    7.890625
#      8.51953125  8.99609375  9.29296875]
#    [ 7.70703125  8.00390625  8.48046875  9.109375    9.515625
#     10.14453125 10.62109375 10.91796875]
#    [10.22265625 10.51953125 10.99609375 11.625      12.03125
#     12.66015625 13.13671875 13.43359375]
#    [12.12890625 12.42578125 12.90234375 13.53125    13.9375
#     14.56640625 15.04296875 15.33984375]
#    [13.31640625 13.61328125 14.08984375 14.71875    15.125
#     15.75390625 16.23046875 16.52734375]]]]
output = interpolate_nd(
    data, lambda x, _: cubic_coeffs(x), scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_cubic",
)

resize_upsample_scales_cubic_A_n0p5_exclude_outside

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
    cubic_coeff_a=-0.5,
    exclude_outside=True,
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)

# [[[[ 0.55882353  0.81494204  1.35698249  1.89705882  2.39705882
#      2.93713516  3.47917561  3.73529412]
#    [ 1.58329755  1.83941606  2.38145651  2.92153285  3.42153285
#      3.96160918  4.50364964  4.75976814]
#    [ 3.75145936  4.00757787  4.54961832  5.08969466  5.58969466
#      6.12977099  6.67181144  6.92792995]
#    [ 5.91176471  6.16788321  6.70992366  7.25        7.75
#      8.29007634  8.83211679  9.08823529]
#    [ 7.91176471  8.16788321  8.70992366  9.25        9.75
#     10.29007634 10.83211679 11.08823529]
#    [10.07207005 10.32818856 10.87022901 11.41030534 11.91030534
#     12.45038168 12.99242213 13.24854064]
#    [12.24023186 12.49635036 13.03839082 13.57846715 14.07846715
#     14.61854349 15.16058394 15.41670245]
#    [13.26470588 13.52082439 14.06286484 14.60294118 15.10294118
#     15.64301751 16.18505796 16.44117647]]]]
output = interpolate_nd(
    data,
    lambda x, _: cubic_coeffs(x, A=-0.5),
    scale_factors=scales,
    exclude_outside=True,
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_cubic_A_n0p5_exclude_outside",
)

resize_upsample_scales_cubic_align_corners

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
    coordinate_transformation_mode="align_corners",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)

# [[[[ 1.          1.34110787  1.80029155  2.32944606  2.67055394
#      3.19970845  3.65889213  4.        ]
#    [ 2.36443149  2.70553936  3.16472303  3.69387755  4.03498542
#      4.56413994  5.02332362  5.36443149]
#    [ 4.20116618  4.54227405  5.00145773  5.53061224  5.87172012
#      6.40087464  6.86005831  7.20116618]
#    [ 6.31778426  6.65889213  7.1180758   7.64723032  7.98833819
#      8.51749271  8.97667638  9.31778426]
#    [ 7.68221574  8.02332362  8.48250729  9.01166181  9.35276968
#      9.8819242  10.34110787 10.68221574]
#    [ 9.79883382 10.13994169 10.59912536 11.12827988 11.46938776
#     11.99854227 12.45772595 12.79883382]
#    [11.63556851 11.97667638 12.43586006 12.96501458 13.30612245
#     13.83527697 14.29446064 14.63556851]
#    [13.         13.34110787 13.80029155 14.32944606 14.67055394
#     15.19970845 15.65889213 16.        ]]]]
output = interpolate_nd(
    data,
    lambda x, _: cubic_coeffs(x),
    scale_factors=scales,
    coordinate_transformation_mode="align_corners",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_cubic_align_corners",
)

resize_upsample_scales_cubic_asymmetric

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="cubic",
    coordinate_transformation_mode="asymmetric",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)

# [[[[ 1.       1.40625  2.       2.5      3.       3.59375  4.
#      4.09375]
#    [ 2.625    3.03125  3.625    4.125    4.625    5.21875  5.625
#      5.71875]
#    [ 5.       5.40625  6.       6.5      7.       7.59375  8.
#      8.09375]
#    [ 7.       7.40625  8.       8.5      9.       9.59375 10.
#     10.09375]
#    [ 9.       9.40625 10.      10.5     11.      11.59375 12.
#     12.09375]
#    [11.375   11.78125 12.375   12.875   13.375   13.96875 14.375
#     14.46875]
#    [13.      13.40625 14.      14.5     15.      15.59375 16.
#     16.09375]
#    [13.375   13.78125 14.375   14.875   15.375   15.96875 16.375
#     16.46875]]]]
output = interpolate_nd(
    data,
    lambda x, _: cubic_coeffs(x, A=-0.75),
    scale_factors=scales,
    coordinate_transformation_mode="asymmetric",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_cubic_asymmetric",
)

resize_upsample_scales_linear

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)

# [[[[1.   1.25 1.75 2.  ]
#    [1.5  1.75 2.25 2.5 ]
#    [2.5  2.75 3.25 3.5 ]
#    [3.   3.25 3.75 4.  ]]]]
output = interpolate_nd(
    data, lambda x, _: linear_coeffs(x), scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_linear",
)

resize_upsample_scales_linear_align_corners

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="align_corners",
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32)

# [[[[1.         1.33333333 1.66666667 2.        ]
#    [1.66666667 2.         2.33333333 2.66666667]
#    [2.33333333 2.66666667 3.         3.33333333]
#    [3.         3.33333333 3.66666667 4.        ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    scale_factors=scales,
    coordinate_transformation_mode="align_corners",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_linear_align_corners",
)

resize_upsample_scales_linear_half_pixel_symmetric

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="linear",
    coordinate_transformation_mode="half_pixel_symmetric",
)

data = np.array([[[[1, 2], [3, 4]]]], dtype=np.float32)
scales = np.array([1.0, 1.0, 2.3, 2.94], dtype=np.float32)

# [[[[1.        , 1.15986395, 1.5       , 1.84013605, 2.        ],
#    [1.56521738, 1.72508133, 2.06521738, 2.40535343, 2.56521738],
#    [2.43478262, 2.59464657, 2.93478262, 3.27491867, 3.43478262],
#    [3.        , 3.15986395, 3.5       , 3.84013605, 4.        ]]]]
output = interpolate_nd(
    data,
    lambda x, _: linear_coeffs(x),
    scale_factors=scales,
    coordinate_transformation_mode="half_pixel_symmetric",
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_linear_half_pixel_symmetric",
)

resize_upsample_scales_nearest

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="nearest",
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 3.0], dtype=np.float32)

# [[[[1. 1. 1. 2. 2. 2.]
#    [1. 1. 1. 2. 2. 2.]
#    [3. 3. 3. 4. 4. 4.]
#    [3. 3. 3. 4. 4. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), scale_factors=scales
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_nearest",
)

resize_upsample_scales_nearest_axes_2_3

axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([2.0, 3.0], dtype=np.float32)

# [[[[1. 1. 1. 2. 2. 2.]
#    [1. 1. 1. 2. 2. 2.]
#    [3. 3. 3. 4. 4. 4.]
#    [3. 3. 3. 4. 4. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), scale_factors=scales, axes=axes
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_nearest_axes_2_3",
)

resize_upsample_scales_nearest_axes_3_2

axes = [3, 2]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "scales"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([3.0, 2.0], dtype=np.float32)

# [[[[1. 1. 1. 2. 2. 2.]
#    [1. 1. 1. 2. 2. 2.]
#    [3. 3. 3. 4. 4. 4.]
#    [3. 3. 3. 4. 4. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), scale_factors=scales, axes=axes
).astype(np.float32)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_resize_upsample_scales_nearest_axes_3_2",
)

resize_upsample_sizes_cubic

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="cubic",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 9, 10], dtype=np.int64)

# [[[[ 0.45507922  0.64057922  0.97157922  1.42257922  1.90732922
#      2.22332922  2.70807922  3.15907922  3.49007922  3.67557922]
#    [ 1.39437963  1.57987963  1.91087963  2.36187963  2.84662963
#      3.16262963  3.64737963  4.09837963  4.42937963  4.61487963]
#    [ 2.95130693  3.13680693  3.46780693  3.91880693  4.40355693
#      4.71955693  5.20430693  5.65530693  5.98630693  6.17180693]
#    [ 5.20525069  5.39075069  5.72175069  6.17275069  6.65750069
#      6.97350069  7.45825069  7.90925069  8.24025069  8.42575069]
#    [ 6.88975     7.07525     7.40625     7.85725     8.342
#      8.658       9.14275     9.59375     9.92475    10.11025   ]
#    [ 8.57424931  8.75974931  9.09074931  9.54174931 10.02649931
#     10.34249931 10.82724931 11.27824931 11.60924931 11.79474931]
#    [10.82819307 11.01369307 11.34469307 11.79569307 12.28044307
#     12.59644307 13.08119307 13.53219307 13.86319307 14.04869307]
#    [12.38512037 12.57062037 12.90162037 13.35262037 13.83737037
#     14.15337037 14.63812037 15.08912037 15.42012037 15.60562037]
#    [13.32442078 13.50992078 13.84092078 14.29192078 14.77667078
#     15.09267078 15.57742078 16.02842078 16.35942078 16.54492078]]]]
output = interpolate_nd(
    data, lambda x, _: cubic_coeffs(x), output_size=sizes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_cubic",
)

resize_upsample_sizes_nearest

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 7, 8], dtype=np.int64)

# [[[[1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), output_size=sizes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest",
)

resize_upsample_sizes_nearest_axes_2_3

axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([7, 8], dtype=np.int64)

# [[[[1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), output_size=sizes, axes=axes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_axes_2_3",
)

resize_upsample_sizes_nearest_axes_3_2

axes = [3, 2]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([8, 7], dtype=np.int64)

# [[[[1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x), output_size=sizes, axes=axes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_axes_3_2",
)

resize_upsample_sizes_nearest_ceil_half_pixel

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    coordinate_transformation_mode="half_pixel",
    nearest_mode="ceil",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 8, 8], dtype=np.int64)

# [[[[ 1.  2.  2.  3.  3.  4.  4.  4.]
#    [ 5.  6.  6.  7.  7.  8.  8.  8.]
#    [ 5.  6.  6.  7.  7.  8.  8.  8.]
#    [ 9. 10. 10. 11. 11. 12. 12. 12.]
#    [ 9. 10. 10. 11. 11. 12. 12. 12.]
#    [13. 14. 14. 15. 15. 16. 16. 16.]
#    [13. 14. 14. 15. 15. 16. 16. 16.]
#    [13. 14. 14. 15. 15. 16. 16. 16.]]]]
output = interpolate_nd(
    data, lambda x, _: nearest_coeffs(x, mode="ceil"), output_size=sizes
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_ceil_half_pixel",
)

resize_upsample_sizes_nearest_floor_align_corners

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    coordinate_transformation_mode="align_corners",
    nearest_mode="floor",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 8, 8], dtype=np.int64)

# [[[[ 1.  1.  1.  2.  2.  3.  3.  4.]
#    [ 1.  1.  1.  2.  2.  3.  3.  4.]
#    [ 1.  1.  1.  2.  2.  3.  3.  4.]
#    [ 5.  5.  5.  6.  6.  7.  7.  8.]
#    [ 5.  5.  5.  6.  6.  7.  7.  8.]
#    [ 9.  9.  9. 10. 10. 11. 11. 12.]
#    [ 9.  9.  9. 10. 10. 11. 11. 12.]
#    [13. 13. 13. 14. 14. 15. 15. 16.]]]]
output = interpolate_nd(
    data,
    lambda x, _: nearest_coeffs(x, mode="floor"),
    output_size=sizes,
    coordinate_transformation_mode="align_corners",
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_floor_align_corners",
)

resize_upsample_sizes_nearest_not_larger

keep_aspect_ratio_policy = "not_larger"
axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([7, 8], dtype=np.int64)  # Results in 7x7

# [[[[1. 1. 1. 1. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2.]
#    [3. 3. 3. 3. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4.]]]]
output = interpolate_nd(
    data,
    lambda x, _: nearest_coeffs(x),
    output_size=sizes,
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_not_larger",
)

resize_upsample_sizes_nearest_not_smaller

keep_aspect_ratio_policy = "not_smaller"
axes = [2, 3]
node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([7, 8], dtype=np.int64)  # Results in 8x8

# [[[[1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [1. 1. 1. 1. 2. 2. 2. 2.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]
#    [3. 3. 3. 3. 4. 4. 4. 4.]]]]
output = interpolate_nd(
    data,
    lambda x, _: nearest_coeffs(x),
    output_size=sizes,
    axes=axes,
    keep_aspect_ratio_policy=keep_aspect_ratio_policy,
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_not_smaller",
)

resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric

node = onnx.helper.make_node(
    "Resize",
    inputs=["X", "", "", "sizes"],
    outputs=["Y"],
    mode="nearest",
    coordinate_transformation_mode="asymmetric",
    nearest_mode="round_prefer_ceil",
)

data = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ],
    dtype=np.float32,
)

sizes = np.array([1, 1, 8, 8], dtype=np.int64)

# [[[[ 1.  2.  2.  3.  3.  4.  4.  4.]
#    [ 5.  6.  6.  7.  7.  8.  8.  8.]
#    [ 5.  6.  6.  7.  7.  8.  8.  8.]
#    [ 9. 10. 10. 11. 11. 12. 12. 12.]
#    [ 9. 10. 10. 11. 11. 12. 12. 12.]
#    [13. 14. 14. 15. 15. 16. 16. 16.]
#    [13. 14. 14. 15. 15. 16. 16. 16.]
#    [13. 14. 14. 15. 15. 16. 16. 16.]]]]
output = interpolate_nd(
    data,
    lambda x, _: nearest_coeffs(x, mode="round_prefer_ceil"),
    output_size=sizes,
    coordinate_transformation_mode="asymmetric",
).astype(np.float32)

expect(
    node,
    inputs=[data, sizes],
    outputs=[output],
    name="test_resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric",
)

ReverseSequence

There are 2 test cases, listed as following:

reversesequence_batch

node = onnx.helper.make_node(
    "ReverseSequence",
    inputs=["x", "sequence_lens"],
    outputs=["y"],
    time_axis=1,
    batch_axis=0,
)
x = np.array(
    [
        [0.0, 1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0, 7.0],
        [8.0, 9.0, 10.0, 11.0],
        [12.0, 13.0, 14.0, 15.0],
    ],
    dtype=np.float32,
)
sequence_lens = np.array([1, 2, 3, 4], dtype=np.int64)

y = np.array(
    [
        [0.0, 1.0, 2.0, 3.0],
        [5.0, 4.0, 6.0, 7.0],
        [10.0, 9.0, 8.0, 11.0],
        [15.0, 14.0, 13.0, 12.0],
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[x, sequence_lens],
    outputs=[y],
    name="test_reversesequence_batch",
)

reversesequence_time

node = onnx.helper.make_node(
    "ReverseSequence",
    inputs=["x", "sequence_lens"],
    outputs=["y"],
    time_axis=0,
    batch_axis=1,
)
x = np.array(
    [
        [0.0, 4.0, 8.0, 12.0],
        [1.0, 5.0, 9.0, 13.0],
        [2.0, 6.0, 10.0, 14.0],
        [3.0, 7.0, 11.0, 15.0],
    ],
    dtype=np.float32,
)
sequence_lens = np.array([4, 3, 2, 1], dtype=np.int64)

y = np.array(
    [
        [3.0, 6.0, 9.0, 12.0],
        [2.0, 5.0, 8.0, 13.0],
        [1.0, 4.0, 10.0, 14.0],
        [0.0, 7.0, 11.0, 15.0],
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[x, sequence_lens],
    outputs=[y],
    name="test_reversesequence_time",
)

RoiAlign

There are 3 test cases, listed as following:

roialign_aligned_false

node = onnx.helper.make_node(
    "RoiAlign",
    inputs=["X", "rois", "batch_indices"],
    outputs=["Y"],
    spatial_scale=1.0,
    output_height=5,
    output_width=5,
    sampling_ratio=2,
    coordinate_transformation_mode="output_half_pixel",
)

X, batch_indices, rois = get_roi_align_input_values()
# (num_rois, C, output_height, output_width)
Y = np.array(
    [
        [
            [
                [0.4664, 0.4466, 0.3405, 0.5688, 0.6068],
                [0.3714, 0.4296, 0.3835, 0.5562, 0.3510],
                [0.2768, 0.4883, 0.5222, 0.5528, 0.4171],
                [0.4713, 0.4844, 0.6904, 0.4920, 0.8774],
                [0.6239, 0.7125, 0.6289, 0.3355, 0.3495],
            ]
        ],
        [
            [
                [0.3022, 0.4305, 0.4696, 0.3978, 0.5423],
                [0.3656, 0.7050, 0.5165, 0.3172, 0.7015],
                [0.2912, 0.5059, 0.6476, 0.6235, 0.8299],
                [0.5916, 0.7389, 0.7048, 0.8372, 0.8893],
                [0.6227, 0.6153, 0.7097, 0.6154, 0.4585],
            ]
        ],
        [
            [
                [0.2384, 0.3379, 0.3717, 0.6100, 0.7601],
                [0.3767, 0.3785, 0.7147, 0.9243, 0.9727],
                [0.5749, 0.5826, 0.5709, 0.7619, 0.8770],
                [0.5355, 0.2566, 0.2141, 0.2796, 0.3600],
                [0.4365, 0.3504, 0.2887, 0.3661, 0.2349],
            ]
        ],
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, rois, batch_indices],
    outputs=[Y],
    name="test_roialign_aligned_false",
)

roialign_aligned_true

node = onnx.helper.make_node(
    "RoiAlign",
    inputs=["X", "rois", "batch_indices"],
    outputs=["Y"],
    spatial_scale=1.0,
    output_height=5,
    output_width=5,
    sampling_ratio=2,
    coordinate_transformation_mode="half_pixel",
)

X, batch_indices, rois = get_roi_align_input_values()
# (num_rois, C, output_height, output_width)
Y = np.array(
    [
        [
            [
                [0.5178, 0.3434, 0.3229, 0.4474, 0.6344],
                [0.4031, 0.5366, 0.4428, 0.4861, 0.4023],
                [0.2512, 0.4002, 0.5155, 0.6954, 0.3465],
                [0.3350, 0.4601, 0.5881, 0.3439, 0.6849],
                [0.4932, 0.7141, 0.8217, 0.4719, 0.4039],
            ]
        ],
        [
            [
                [0.3070, 0.2187, 0.3337, 0.4880, 0.4870],
                [0.1871, 0.4914, 0.5561, 0.4192, 0.3686],
                [0.1433, 0.4608, 0.5971, 0.5310, 0.4982],
                [0.2788, 0.4386, 0.6022, 0.7000, 0.7524],
                [0.5774, 0.7024, 0.7251, 0.7338, 0.8163],
            ]
        ],
        [
            [
                [0.2393, 0.4075, 0.3379, 0.2525, 0.4743],
                [0.3671, 0.2702, 0.4105, 0.6419, 0.8308],
                [0.5556, 0.4543, 0.5564, 0.7502, 0.9300],
                [0.6626, 0.5617, 0.4813, 0.4954, 0.6663],
                [0.6636, 0.3721, 0.2056, 0.1928, 0.2478],
            ]
        ],
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[X, rois, batch_indices],
    outputs=[Y],
    name="test_roialign_aligned_true",
)

roialign_mode_max

X = np.array(
    [
        [
            [
                [
                    0.2764,
                    0.715,
                    0.1958,
                    0.3416,
                    0.4638,
                    0.0259,
                    0.2963,
                    0.6518,
                    0.4856,
                    0.725,
                ],
                [
                    0.9637,
                    0.0895,
                    0.2919,
                    0.6753,
                    0.0234,
                    0.6132,
                    0.8085,
                    0.5324,
                    0.8992,
                    0.4467,
                ],
                [
                    0.3265,
                    0.8479,
                    0.9698,
                    0.2471,
                    0.9336,
                    0.1878,
                    0.4766,
                    0.4308,
                    0.34,
                    0.2162,
                ],
                [
                    0.0206,
                    0.172,
                    0.2155,
                    0.4394,
                    0.0653,
                    0.3406,
                    0.7724,
                    0.3921,
                    0.2541,
                    0.5799,
                ],
                [
                    0.4062,
                    0.2194,
                    0.4473,
                    0.4687,
                    0.7109,
                    0.9327,
                    0.9815,
                    0.632,
                    0.1728,
                    0.6119,
                ],
                [
                    0.3097,
                    0.1283,
                    0.4984,
                    0.5068,
                    0.4279,
                    0.0173,
                    0.4388,
                    0.043,
                    0.4671,
                    0.7119,
                ],
                [
                    0.1011,
                    0.8477,
                    0.4726,
                    0.1777,
                    0.9923,
                    0.4042,
                    0.1869,
                    0.7795,
                    0.9946,
                    0.9689,
                ],
                [
                    0.1366,
                    0.3671,
                    0.7011,
                    0.6234,
                    0.9867,
                    0.5585,
                    0.6985,
                    0.5609,
                    0.8788,
                    0.9928,
                ],
                [
                    0.5697,
                    0.8511,
                    0.6711,
                    0.9406,
                    0.8751,
                    0.7496,
                    0.165,
                    0.1049,
                    0.1559,
                    0.2514,
                ],
                [
                    0.7012,
                    0.4056,
                    0.7879,
                    0.3461,
                    0.0415,
                    0.2998,
                    0.5094,
                    0.3727,
                    0.5482,
                    0.0502,
                ],
            ]
        ]
    ],
    dtype=np.float32,
)
rois = np.array(
    [[0.0, 0.0, 9.0, 9.0], [0.0, 5.0, 4.0, 9.0], [5.0, 5.0, 9.0, 9.0]],
    dtype=np.float32,
)
batch_indices = np.array([0, 0, 0], dtype=np.int64)

Y = np.array(
    [
        [
            [
                [0.3445228, 0.37310338, 0.37865096, 0.446696, 0.37991184],
                [0.4133513, 0.5455125, 0.6651902, 0.55805874, 0.27110294],
                [0.21223956, 0.40924096, 0.8417618, 0.792561, 0.37196714],
                [0.46835402, 0.39741728, 0.8012819, 0.4969306, 0.5495158],
                [0.3595896, 0.5196813, 0.5403741, 0.23814403, 0.19992709],
            ]
        ],
        [
            [
                [0.30517197, 0.5086199, 0.3189761, 0.4054401, 0.47630402],
                [0.50862, 0.8477, 0.37808004, 0.24936005, 0.79384017],
                [0.17620805, 0.29368007, 0.44870415, 0.4987201, 0.63148826],
                [0.51066005, 0.8511, 0.5368801, 0.9406, 0.70008016],
                [0.4487681, 0.51066035, 0.5042561, 0.5643603, 0.42004836],
            ]
        ],
        [
            [
                [0.21062402, 0.3510401, 0.37416005, 0.5967599, 0.46507207],
                [0.32336006, 0.31180006, 0.6236001, 0.9946, 0.7751202],
                [0.35744014, 0.5588001, 0.35897616, 0.7030401, 0.6353923],
                [0.5996801, 0.27940005, 0.17948808, 0.35152006, 0.31769615],
                [0.3598083, 0.40752012, 0.2385281, 0.43856013, 0.26313624],
            ]
        ],
    ],
    dtype=np.float32,
)

node = onnx.helper.make_node(
    "RoiAlign",
    inputs=["X", "rois", "batch_indices"],
    mode="max",
    outputs=["Y"],
    spatial_scale=1.0,
    output_height=5,
    output_width=5,
    sampling_ratio=2,
    coordinate_transformation_mode="output_half_pixel",
)

expect(
    node,
    inputs=[X, rois, batch_indices],
    outputs=[Y],
    name="test_roialign_mode_max",
)

RotaryEmbedding

There are 8 test cases, listed as following:

rotary_embedding

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache", "position_ids"],
    outputs=["output"],
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
sin_cache_data = np.random.rand(50, 4).astype(np.float32)
cos_cache_data = np.random.rand(50, 4).astype(np.float32)

expected_output = rotary_embedding(
    input_data, cos_cache_data, sin_cache_data, position_ids=position_ids_data
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
    outputs=[expected_output],
    name="test_rotary_embedding",
)

rotary_embedding_3d_input

num_heads = 4
node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache", "position_ids"],
    outputs=["output"],
    num_heads=num_heads,
)

input_data = np.random.rand(2, 3, 32).astype(np.float32)
position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
sin_cache_data = np.random.rand(50, 4).astype(np.float32)
cos_cache_data = np.random.rand(50, 4).astype(np.float32)

expected_output = rotary_embedding(
    input_data,
    cos_cache_data,
    sin_cache_data,
    position_ids=position_ids_data,
    num_heads=num_heads,
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
    outputs=[expected_output],
    name="test_rotary_embedding_3d_input",
)

rotary_embedding_interleaved

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache", "position_ids"],
    outputs=["output"],
    interleaved=1,
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
sin_cache_data = np.random.rand(50, 4).astype(np.float32)
cos_cache_data = np.random.rand(50, 4).astype(np.float32)

expected_output = rotary_embedding(
    input_data,
    cos_cache_data,
    sin_cache_data,
    position_ids=position_ids_data,
    interleaved=1,
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
    outputs=[expected_output],
    name="test_rotary_embedding_interleaved",
)

rotary_embedding_no_position_ids

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache"],
    outputs=["output"],
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
sin_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
cos_cache_data = np.random.rand(2, 3, 4).astype(np.float32)

expected_output = rotary_embedding(input_data, cos_cache_data, sin_cache_data)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data],
    outputs=[expected_output],
    name="test_rotary_embedding_no_position_ids",
)

rotary_embedding_no_position_ids_interleaved

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache"],
    outputs=["output"],
    interleaved=1,
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
sin_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
cos_cache_data = np.random.rand(2, 3, 4).astype(np.float32)

expected_output = rotary_embedding(
    input_data,
    cos_cache_data,
    sin_cache_data,
    interleaved=1,
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data],
    outputs=[expected_output],
    name="test_rotary_embedding_no_position_ids_interleaved",
)

rotary_embedding_no_position_ids_rotary_dim

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache"],
    outputs=["output"],
    rotary_embedding_dim=4,
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
sin_cache_data = np.random.rand(2, 3, 2).astype(np.float32)
cos_cache_data = np.random.rand(2, 3, 2).astype(np.float32)

expected_output = rotary_embedding(
    input_data,
    cos_cache_data,
    sin_cache_data,
    rotary_embedding_dim=4,
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data],
    outputs=[expected_output],
    name="test_rotary_embedding_no_position_ids_rotary_dim",
)

rotary_embedding_with_interleaved_rotary_dim

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache", "position_ids"],
    outputs=["output"],
    rotary_embedding_dim=4,
    interleaved=1,
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
sin_cache_data = np.random.rand(50, 2).astype(np.float32)
cos_cache_data = np.random.rand(50, 2).astype(np.float32)

expected_output = rotary_embedding(
    input_data,
    cos_cache_data,
    sin_cache_data,
    position_ids=position_ids_data,
    interleaved=1,
    rotary_embedding_dim=4,
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
    outputs=[expected_output],
    name="test_rotary_embedding_with_interleaved_rotary_dim",
)

rotary_embedding_with_rotary_dim

node = onnx.helper.make_node(
    "RotaryEmbedding",
    inputs=["input", "cos_cache", "sin_cache", "position_ids"],
    outputs=["output"],
    rotary_embedding_dim=4,
)

input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
sin_cache_data = np.random.rand(50, 2).astype(np.float32)
cos_cache_data = np.random.rand(50, 2).astype(np.float32)

expected_output = rotary_embedding(
    input_data,
    cos_cache_data,
    sin_cache_data,
    position_ids=position_ids_data,
    rotary_embedding_dim=4,
)

expect(
    node,
    inputs=[input_data, cos_cache_data, sin_cache_data, position_ids_data],
    outputs=[expected_output],
    name="test_rotary_embedding_with_rotary_dim",
)

Round

There are 1 test cases, listed as following:

round

node = onnx.helper.make_node(
    "Round",
    inputs=["x"],
    outputs=["y"],
)

x = np.array(
    [
        0.1,
        0.5,
        0.9,
        1.2,
        1.5,
        1.8,
        2.3,
        2.5,
        2.7,
        -1.1,
        -1.5,
        -1.9,
        -2.2,
        -2.5,
        -2.8,
    ]
).astype(np.float32)

# expected output
y = np.array(
    [
        0.0,
        0.0,
        1.0,
        1.0,
        2.0,
        2.0,
        2.0,
        2.0,
        3.0,
        -1.0,
        -2.0,
        -2.0,
        -2.0,
        -2.0,
        -3.0,
    ]
).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_round")

STFT

There are 1 test cases, listed as following:

stft

signal = np.arange(0, 128, dtype=np.float32).reshape(1, 128, 1)
length = np.array(16).astype(np.int64)
onesided_length = (length >> 1) + 1
step = np.array(8).astype(np.int64)

no_window = ""  # optional input, not supplied
node = onnx.helper.make_node(
    "STFT",
    inputs=["signal", "frame_step", no_window, "frame_length"],
    outputs=["output"],
)

nstfts = ((signal.shape[1] - length) // step) + 1
# [batch_size][frames][frame_length][2]
output = np.empty([1, nstfts, onesided_length, 2], dtype=np.float32)
for i in range(nstfts):
    start = i * step
    stop = i * step + length
    complex_out = np.fft.fft(signal[0, start:stop, 0])[0:onesided_length]
    output[0, i] = np.stack((complex_out.real, complex_out.imag), axis=1)

output = output.astype(signal.dtype)
expect(node, inputs=[signal, step, length], outputs=[output], name="test_stft")

node = onnx.helper.make_node(
    "STFT",
    inputs=["signal", "frame_step", "window"],
    outputs=["output"],
)

# Test with window
a0 = 0.5
a1 = 0.5
window = a0 + a1 * np.cos(
    2 * np.pi * np.arange(0, length, 1, dtype=np.float32) / length
)
nstfts = 1 + (signal.shape[1] - window.shape[0]) // step

# [batch_size][frames][frame_length][2]
output = np.empty([1, nstfts, onesided_length, 2], dtype=np.float32)
for i in range(nstfts):
    start = i * step
    stop = i * step + length
    complex_out = np.fft.fft(signal[0, start:stop, 0] * window)[
        0:onesided_length
    ]
    output[0, i] = np.stack((complex_out.real, complex_out.imag), axis=1)
window = window.astype(signal.dtype)
output = output.astype(signal.dtype)
expect(
    node,
    inputs=[signal, step, window],
    outputs=[output],
    name="test_stft_with_window",
)

Scan

There are 4 test cases, listed as following:

scan_8

# Given an input sequence [x1, ..., xN], sum up its elements using a scan
# returning the final state (x1+x2+...+xN) as well the scan_output
# [x1, x1+x2, ..., x1+x2+...+xN]
# Note: the first input (sequence_lens) is optional and omitted via "".
node = onnx.parser.parse_node(
    """
    y, z = Scan ("", initial, x) <
        num_scan_inputs = 1,
        body = scan_body (float[2] sum_in, float[2] next)
            => (float[2] sum_out, float[2] scan_out)
        {
            sum_out  = Add(sum_in, next)
            scan_out = Identity(sum_out)
        }
    >
    """
)
# create inputs for batch-size 1, sequence-length 3, inner dimension 2
initial = np.array([0, 0]).astype(np.float32).reshape((1, 2))
x = np.array([1, 2, 3, 4, 5, 6]).astype(np.float32).reshape((1, 3, 2))
# final state computed = [1 + 3 + 5, 2 + 4 + 6]
y = np.array([9, 12]).astype(np.float32).reshape((1, 2))
# scan-output computed
z = np.array([1, 2, 4, 6, 9, 12]).astype(np.float32).reshape((1, 3, 2))

expect(
    node,
    inputs=[initial, x],
    outputs=[y, z],
    name="test_scan_sum",
    opset_imports=[onnx.helper.make_opsetid("", 8)],
)

scan_9

# Given an input sequence [x1, ..., xN], sum up its elements using a scan
# returning the final state (x1+x2+...+xN) as well the scan_output
# [x1, x1+x2, ..., x1+x2+...+xN]
node = onnx.parser.parse_node(
    """
    y, z = Scan (initial, x) <
        num_scan_inputs = 1,
        body = scan_body (float[2] sum_in, float[2] next)
            => (float[2] sum_out, float[2] scan_out)
        {
            sum_out  = Add(sum_in, next)
            scan_out = Identity(sum_out)
        }
    >
    """
)
# create inputs for sequence-length 3, inner dimension 2
initial = np.array([0, 0]).astype(np.float32).reshape((2,))
x = np.array([1, 2, 3, 4, 5, 6]).astype(np.float32).reshape((3, 2))
# final state computed = [1 + 3 + 5, 2 + 4 + 6]
y = np.array([9, 12]).astype(np.float32).reshape((2,))
# scan-output computed
z = np.array([1, 2, 4, 6, 9, 12]).astype(np.float32).reshape((3, 2))

expect(
    node,
    inputs=[initial, x],
    outputs=[y, z],
    name="test_scan9_sum",
    opset_imports=[onnx.helper.make_opsetid("", 9)],
)

scan_9_multi_state

# Scan with two state variables: running sum and running product.
# This exercises the case where num_loop_state_vars (2) differs from
# num_scan_inputs (1).
#
# Body inputs:  sum_in (state), prod_in (state), next (scan)
# Body outputs: sum_out (state), prod_out (state), scan_out (scan)
node = onnx.parser.parse_node(
    """
    y_sum, y_prod, z = Scan (initial_sum, initial_prod, x) <
        num_scan_inputs = 1,
        body = scan_body (float[2] sum_in, float[2] prod_in, float[2] next)
            => (float[2] sum_out, float[2] prod_out, float[2] scan_out)
        {
            sum_out  = Add(sum_in, next)
            prod_out = Mul(prod_in, next)
            scan_out = Identity(sum_out)
        }
    >
    """
)
# x = [[1, 2], [3, 4], [5, 6]]
initial_sum = np.array([0, 0]).astype(np.float32)
initial_prod = np.array([1, 1]).astype(np.float32)
x = np.array([1, 2, 3, 4, 5, 6]).astype(np.float32).reshape((3, 2))
# final sum = [1+3+5, 2+4+6] = [9, 12]
y_sum = np.array([9, 12]).astype(np.float32)
# final product = [1*3*5, 2*4*6] = [15, 48]
y_prod = np.array([15, 48]).astype(np.float32)
# scan output (running sum) = [[1,2], [4,6], [9,12]]
z = np.array([1, 2, 4, 6, 9, 12]).astype(np.float32).reshape((3, 2))

expect(
    node,
    inputs=[initial_sum, initial_prod, x],
    outputs=[y_sum, y_prod, z],
    name="test_scan9_multi_state",
    opset_imports=[onnx.helper.make_opsetid("", 9)],
)

scan_9_scalar

# Scan with scalar state and scan output to verify that output
# shapes are not distorted (e.g. (T,) not (T, 1)).
node = onnx.parser.parse_node(
    """
    y, z = Scan (initial, x) <
        num_scan_inputs = 1,
        body = scan_body (float sum_in, float next)
            => (float sum_out, float scan_out)
        {
            sum_out  = Add(sum_in, next)
            scan_out = Identity(sum_out)
        }
    >
    """
)
initial = np.float32(0.0)
x = np.array([1, 2, 3, 4, 5]).astype(np.float32)
# final state = 1+2+3+4+5 = 15
y = np.float32(15.0)
# scan output = [1, 3, 6, 10, 15], shape (5,)
z = np.array([1, 3, 6, 10, 15]).astype(np.float32)

expect(
    node,
    inputs=[initial, x],
    outputs=[y, z],
    name="test_scan9_scalar",
    opset_imports=[onnx.helper.make_opsetid("", 9)],
)

Scatter

There are 2 test cases, listed as following:

scatter_with_axis

axis = 1
node = onnx.helper.make_node(
    "Scatter",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, 3]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter(data, indices, updates, axis=axis)
# print(y) produces
# [[1.0, 1.1, 3.0, 2.1, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_with_axis",
    opset_imports=[helper.make_opsetid("", 10)],
)

scatter_without_axis

node = onnx.helper.make_node(
    "Scatter",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
)
data = np.zeros((3, 3), dtype=np.float32)
indices = np.array([[1, 0, 2], [0, 2, 1]], dtype=np.int64)
updates = np.array([[1.0, 1.1, 1.2], [2.0, 2.1, 2.2]], dtype=np.float32)

y = scatter(data, indices, updates)
# print(y) produces
# [[2.0, 1.1, 0.0],
#  [1.0, 0.0, 2.2],
#  [0.0, 2.1, 1.2]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_without_axis",
    opset_imports=[helper.make_opsetid("", 10)],
)

ScatterElements

There are 7 test cases, listed as following:

scatter_elements_with_axis

axis = 1
node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, 3]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter_elements(data, indices, updates, axis)
# print(y) produces
# [[1.0, 1.1, 3.0, 2.1, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_with_axis",
)

scatter_elements_with_duplicate_indices

axis = 1
node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
    reduction="add",
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, 1]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter_elements(data, indices, updates, axis, reduction="add")
# print(y) produces
# [[1.0, 5.2, 3.0, 4.0, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_with_duplicate_indices",
)

scatter_elements_with_negative_indices

axis = 1
node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, -3]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter_elements(data, indices, updates, axis)
# print(y) produces
# [[1.0, 1.1, 2.1, 4.0, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_with_negative_indices",
)

scatter_elements_with_reduction_max

axis = 1
node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
    reduction="max",
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, 1]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter_elements(data, indices, updates, axis, reduction="max")
# print(y) produces
# [[1.0, 2.1, 3.0, 4.0, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_with_reduction_max",
)

scatter_elements_with_reduction_min

axis = 1
node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
    reduction="min",
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, 1]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter_elements(data, indices, updates, axis, reduction="min")
# print(y) produces
# [[1.0, 1.1, 3.0, 4.0, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_with_reduction_min",
)

scatter_elements_with_reduction_mul

axis = 1
node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    axis=axis,
    reduction="mul",
)
data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
indices = np.array([[1, 1]], dtype=np.int64)
updates = np.array([[1.1, 2.1]], dtype=np.float32)

y = scatter_elements(data, indices, updates, axis, reduction="mul")
# print(y) produces
# [[1.0, 4.62, 3.0, 4.0, 5.0]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_with_reduction_mul",
)

scatter_elements_without_axis

node = onnx.helper.make_node(
    "ScatterElements",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
)
data = np.zeros((3, 3), dtype=np.float32)
indices = np.array([[1, 0, 2], [0, 2, 1]], dtype=np.int64)
updates = np.array([[1.0, 1.1, 1.2], [2.0, 2.1, 2.2]], dtype=np.float32)

y = scatter_elements(data, indices, updates)
# print(y) produces
# [[2.0, 1.1, 0.0],
#  [1.0, 0.0, 2.2],
#  [0.0, 2.1, 1.2]]

expect(
    node,
    inputs=[data, indices, updates],
    outputs=[y],
    name="test_scatter_elements_without_axis",
)

ScatterND

There are 7 test cases, listed as following:

scatternd

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
)
data = np.array(
    [
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
    ],
    dtype=np.float32,
)
indices = np.array([[0], [2]], dtype=np.int64)
updates = np.array(
    [
        [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
        [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
    ],
    dtype=np.float32,
)
# Expecting output as np.array(
#    [[[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
#     [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
#     [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
#     [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates)
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd",
)

scatternd_add

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    reduction="add",
)
data = np.array(
    [
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
    ],
    dtype=np.float32,
)
indices = np.array([[0], [0]], dtype=np.int64)
updates = np.array(
    [
        [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
        [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
    ],
    dtype=np.float32,
)
# Expecting output as np.array(
#    [[[7, 8, 9, 10], [13, 14, 15, 16], [18, 17, 16, 15], [16, 15, 14, 13]],
#     [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
#     [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
#     [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates, reduction="add")
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd_add",
)

scatternd_max

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    reduction="max",
)
data = np.array(
    [
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
    ],
    dtype=np.float32,
)
indices = np.array([[0], [0]], dtype=np.int64)
updates = np.array(
    [
        [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
        [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
    ],
    dtype=np.float32,
)
# Expecting output as np.array(
#    [[[5, 5, 5, 5], [6, 6, 7, 8], [8, 7, 7, 7], [8, 8 ,8, 8]],
#     [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
#     [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
#     [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates, reduction="max")
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd_max",
)

scatternd_max_with_element_indices

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    reduction="max",
)
data = np.array([[1, 2], [3, 4]], dtype=np.float32)
# Indices address individual elements (index rank == data rank),
# which exercises the reduction at the element level.
indices = np.array([[0, 0], [1, 1]], dtype=np.int64)
updates = np.array([5, 1], dtype=np.float32)
# Expecting output as np.array([[5, 2], [3, 4]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates, reduction="max")
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd_max_with_element_indices",
)

scatternd_min

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    reduction="min",
)
data = np.array(
    [
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
    ],
    dtype=np.float32,
)
indices = np.array([[0], [0]], dtype=np.int64)
updates = np.array(
    [
        [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
        [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
    ],
    dtype=np.float32,
)
# Expecting output as np.array(
#    [[[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 3, 2, 1]],
#     [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
#     [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
#     [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates, reduction="min")
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd_min",
)

scatternd_min_with_element_indices

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    reduction="min",
)
data = np.array([[1, 2], [3, 4]], dtype=np.float32)
indices = np.array([[0, 0], [1, 1]], dtype=np.int64)
updates = np.array([5, 1], dtype=np.float32)
# Expecting output as np.array([[1, 2], [3, 1]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates, reduction="min")
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd_min_with_element_indices",
)

scatternd_multiply

node = onnx.helper.make_node(
    "ScatterND",
    inputs=["data", "indices", "updates"],
    outputs=["y"],
    reduction="mul",
)
data = np.array(
    [
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
        [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
    ],
    dtype=np.float32,
)
indices = np.array([[0], [0]], dtype=np.int64)
updates = np.array(
    [
        [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
        [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
    ],
    dtype=np.float32,
)
# Expecting output as np.array(
#    [[[5, 10, 15, 20], [60, 72, 84, 96], [168, 147, 126, 105], [128, 96, 64, 32]],
#     [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
#     [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
#     [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]], dtype=np.float32)
output = scatter_nd_impl(data, indices, updates, reduction="mul")
expect(
    node,
    inputs=[data, indices, updates],
    outputs=[output],
    name="test_scatternd_multiply",
)

Selu

There are 2 test cases, listed as following:

selu

node = onnx.helper.make_node(
    "Selu", inputs=["x"], outputs=["y"], alpha=2.0, gamma=3.0
)

x = np.array([-1, 0, 1]).astype(np.float32)
# expected output [-3.79272318, 0., 3.]
y = (
    np.clip(x, 0, np.inf) * 3.0
    + (np.exp(np.clip(x, -np.inf, 0)) - 1) * 2.0 * 3.0
)
expect(node, inputs=[x], outputs=[y], name="test_selu_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = (
    np.clip(x, 0, np.inf) * 3.0
    + (np.exp(np.clip(x, -np.inf, 0)) - 1) * 2.0 * 3.0
)
expect(node, inputs=[x], outputs=[y], name="test_selu")

selu_default

default_alpha = 1.67326319217681884765625
default_gamma = 1.05070102214813232421875
node = onnx.helper.make_node(
    "Selu",
    inputs=["x"],
    outputs=["y"],
)
x = np.random.randn(3, 4, 5).astype(np.float32)
y = (
    np.clip(x, 0, np.inf) * default_gamma
    + (np.exp(np.clip(x, -np.inf, 0)) - 1) * default_alpha * default_gamma
)
expect(node, inputs=[x], outputs=[y], name="test_selu_default")

SequenceInsert

There are 1 test cases, listed as following:

sequenceinsert

test_cases = {
    "at_back": [np.array([10, 11, 12]).astype(np.int64)],
    "at_front": [np.array([-2, -1, 0]), np.array([0]).astype(np.int64)],
}
sequence = [
    np.array([1, 2, 3, 4]).astype(np.int64),
    np.array([5, 6, 7]).astype(np.int64),
    np.array([8, 9]).astype(np.int64),
]

for test_name, test_inputs in test_cases.items():
    tensor = test_inputs[0].astype(np.int64)

    if len(test_inputs) > 1:
        node = onnx.helper.make_node(
            "SequenceInsert",
            inputs=["sequence", "tensor", "position"],
            outputs=["output_sequence"],
        )
        position = test_inputs[1]
        inserted = sequence_insert_reference_implementation(
            sequence, tensor, position
        )
        expect(
            node,
            inputs=[sequence, tensor, position],
            outputs=[inserted],
            name="test_sequence_insert_" + test_name,
        )
    else:
        node = onnx.helper.make_node(
            "SequenceInsert",
            inputs=["sequence", "tensor"],
            outputs=["output_sequence"],
        )
        inserted = sequence_insert_reference_implementation(sequence, tensor)
        expect(
            node,
            inputs=[sequence, tensor],
            outputs=[inserted],
            name="test_sequence_insert_" + test_name,
        )

SequenceMap

There are 6 test cases, listed as following:

sequence_map_add_1_sequence_1_tensor

body = onnx.helper.make_graph(
    [onnx.helper.make_node("Add", ["in0", "in1"], ["out0"])],
    "seq_map_body",
    [
        onnx.helper.make_tensor_value_info(
            "in0", onnx.TensorProto.FLOAT, ["N"]
        ),
        onnx.helper.make_tensor_value_info(
            "in1", onnx.TensorProto.FLOAT, ["N"]
        ),
    ],
    [onnx.helper.make_tensor_value_info("out0", onnx.TensorProto.FLOAT, ["N"])],
)

node = onnx.helper.make_node(
    "SequenceMap", inputs=["x0", "x1"], outputs=["y0"], body=body
)

x0 = [np.random.uniform(0.0, 1.0, 10).astype(np.float32) for k in range(3)]
x1 = np.random.uniform(0.0, 1.0, 10).astype(np.float32)
y0 = [x0[i] + x1 for i in range(3)]
input_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
    onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"]),
]
output_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
]
expect(
    node,
    inputs=[x0, x1],
    outputs=[y0],
    input_type_protos=input_type_protos,
    output_type_protos=output_type_protos,
    name="test_sequence_map_add_1_sequence_1_tensor",
)

sequence_map_add_2_sequences

body = onnx.helper.make_graph(
    [onnx.helper.make_node("Add", ["in0", "in1"], ["out0"])],
    "seq_map_body",
    [
        onnx.helper.make_tensor_value_info(
            "in0", onnx.TensorProto.FLOAT, ["N"]
        ),
        onnx.helper.make_tensor_value_info(
            "in1", onnx.TensorProto.FLOAT, ["N"]
        ),
    ],
    [onnx.helper.make_tensor_value_info("out0", onnx.TensorProto.FLOAT, ["N"])],
)

node = onnx.helper.make_node(
    "SequenceMap", inputs=["x0", "x1"], outputs=["y0"], body=body
)

N = [np.random.randint(1, 10) for _ in range(3)]
x0 = [np.random.uniform(0.0, 1.0, N[k]).astype(np.float32) for k in range(3)]
x1 = [np.random.uniform(0.0, 1.0, N[k]).astype(np.float32) for k in range(3)]
y0 = [x0[k] + x1[k] for k in range(3)]
input_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
]
output_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
]
expect(
    node,
    inputs=[x0, x1],
    outputs=[y0],
    input_type_protos=input_type_protos,
    output_type_protos=output_type_protos,
    name="test_sequence_map_add_2_sequences",
)

sequence_map_extract_shapes

body = onnx.helper.make_graph(
    [onnx.helper.make_node("Shape", ["x"], ["shape"])],
    "seq_map_body",
    [
        onnx.helper.make_tensor_value_info(
            "x", onnx.TensorProto.FLOAT, ["H", "W", "C"]
        )
    ],
    [onnx.helper.make_tensor_value_info("shape", onnx.TensorProto.INT64, [3])],
)

node = onnx.helper.make_node(
    "SequenceMap", inputs=["in_seq"], outputs=["shapes"], body=body
)

shapes = [
    np.array([40, 30, 3], dtype=np.int64),
    np.array([20, 10, 3], dtype=np.int64),
    np.array([10, 5, 3], dtype=np.int64),
]
x0 = [np.zeros(shape, dtype=np.float32) for shape in shapes]
input_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(
            onnx.TensorProto.FLOAT, ["H", "W", "C"]
        )
    ),
]
output_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.INT64, [3])
    ),
]
expect(
    node,
    inputs=[x0],
    outputs=[shapes],
    input_type_protos=input_type_protos,
    output_type_protos=output_type_protos,
    name="test_sequence_map_extract_shapes",
)

sequence_map_identity_1_sequence

body = onnx.helper.make_graph(
    [onnx.helper.make_node("Identity", ["in0"], ["out0"])],
    "seq_map_body",
    [onnx.helper.make_tensor_value_info("in0", onnx.TensorProto.FLOAT, ["N"])],
    [onnx.helper.make_tensor_value_info("out0", onnx.TensorProto.FLOAT, ["M"])],
)

node = onnx.helper.make_node(
    "SequenceMap", inputs=["x"], outputs=["y"], body=body
)

x = [np.random.uniform(0.0, 1.0, 10).astype(np.float32) for _ in range(3)]
y = x
input_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
]
output_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
]
expect(
    node,
    inputs=[x],
    outputs=[y],
    input_type_protos=input_type_protos,
    output_type_protos=output_type_protos,
    name="test_sequence_map_identity_1_sequence",
)

sequence_map_identity_1_sequence_1_tensor

body = onnx.helper.make_graph(
    [
        onnx.helper.make_node("Identity", ["in0"], ["out0"]),
        onnx.helper.make_node("Identity", ["in1"], ["out1"]),
    ],
    "seq_map_body",
    [
        onnx.helper.make_tensor_value_info(
            "in0", onnx.TensorProto.FLOAT, ["N"]
        ),
        onnx.helper.make_tensor_value_info(
            "in1", onnx.TensorProto.FLOAT, ["M"]
        ),
    ],
    [
        onnx.helper.make_tensor_value_info(
            "out0", onnx.TensorProto.FLOAT, ["N"]
        ),
        onnx.helper.make_tensor_value_info(
            "out1", onnx.TensorProto.FLOAT, ["M"]
        ),
    ],
)

node = onnx.helper.make_node(
    "SequenceMap", inputs=["x0", "x1"], outputs=["y0", "y1"], body=body
)

x0 = [
    np.random.uniform(0.0, 1.0, np.random.randint(1, 10)).astype(np.float32)
    for _ in range(3)
]
x1 = np.random.uniform(0.0, 1.0, np.random.randint(1, 10)).astype(np.float32)
y0 = x0
y1 = [x1 for _ in range(3)]
input_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
    onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["M"]),
]
output_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["M"])
    ),
]
expect(
    node,
    inputs=[x0, x1],
    outputs=[y0, y1],
    input_type_protos=input_type_protos,
    output_type_protos=output_type_protos,
    name="test_sequence_map_identity_1_sequence_1_tensor",
)

sequence_map_identity_2_sequences

body = onnx.helper.make_graph(
    [
        onnx.helper.make_node("Identity", ["in0"], ["out0"]),
        onnx.helper.make_node("Identity", ["in1"], ["out1"]),
    ],
    "seq_map_body",
    [
        onnx.helper.make_tensor_value_info(
            "in0", onnx.TensorProto.FLOAT, ["N"]
        ),
        onnx.helper.make_tensor_value_info(
            "in1", onnx.TensorProto.FLOAT, ["M"]
        ),
    ],
    [
        onnx.helper.make_tensor_value_info(
            "out0", onnx.TensorProto.FLOAT, ["N"]
        ),
        onnx.helper.make_tensor_value_info(
            "out1", onnx.TensorProto.FLOAT, ["M"]
        ),
    ],
)

node = onnx.helper.make_node(
    "SequenceMap", inputs=["x0", "x1"], outputs=["y0", "y1"], body=body
)

x0 = [
    np.random.uniform(0.0, 1.0, np.random.randint(1, 10)).astype(np.float32)
    for _ in range(3)
]
x1 = [
    np.random.uniform(0.0, 1.0, np.random.randint(1, 10)).astype(np.float32)
    for _ in range(3)
]
y0 = x0
y1 = x1
input_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["M"])
    ),
]
output_type_protos = [
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["N"])
    ),
    onnx.helper.make_sequence_type_proto(
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, ["M"])
    ),
]
expect(
    node,
    inputs=[x0, x1],
    outputs=[y0, y1],
    input_type_protos=input_type_protos,
    output_type_protos=output_type_protos,
    name="test_sequence_map_identity_2_sequences",
)

Shape

There are 1 test cases, listed as following:

shape

x = np.array(
    [
        [1, 2, 3],
        [4, 5, 6],
    ]
).astype(np.float32)
test_shape("_example", x)  # preserve names of original test cases

x = np.random.randn(3, 4, 5).astype(np.float32)

test_shape("", x)  # preserve names of original test cases

test_shape("_start_1", x, start=1)

test_shape("_end_1", x, end=1)

test_shape("_start_negative_1", x, start=-1)

test_shape("_end_negative_1", x, end=-1)

test_shape("_start_1_end_negative_1", x, start=1, end=-1)

test_shape("_start_1_end_2", x, start=1, end=2)

test_shape("_clip_start", x, start=-10)

test_shape("_clip_end", x, end=10)

test_shape("_start_greater_than_end", x, start=2, end=1)

Shrink

There are 2 test cases, listed as following:

hard_shrink

node = onnx.helper.make_node(
    "Shrink",
    inputs=["x"],
    outputs=["y"],
    lambd=1.5,
)
X = np.arange(-2.0, 2.1, dtype=np.float32)
Y = np.array([-2, 0, 0, 0, 2], dtype=np.float32)
expect(node, inputs=[X], outputs=[Y], name="test_shrink_hard")

soft_shrink

node = onnx.helper.make_node(
    "Shrink",
    inputs=["x"],
    outputs=["y"],
    lambd=1.5,
    bias=1.5,
)
X = np.arange(-2.0, 2.1, dtype=np.float32)
Y = np.array([-0.5, 0, 0, 0, 0.5], dtype=np.float32)
expect(node, inputs=[X], outputs=[Y], name="test_shrink_soft")

Sigmoid

There are 1 test cases, listed as following:

sigmoid

node = onnx.helper.make_node(
    "Sigmoid",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = 1.0 / (
    1.0 + np.exp(np.negative(x))
)  # expected output [0.26894143, 0.5, 0.7310586]
expect(node, inputs=[x], outputs=[y], name="test_sigmoid_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = 1.0 / (1.0 + np.exp(np.negative(x)))
expect(node, inputs=[x], outputs=[y], name="test_sigmoid")

Sign

There are 1 test cases, listed as following:

sign

node = onnx.helper.make_node(
    "Sign",
    inputs=["x"],
    outputs=["y"],
)

x = np.array(range(-5, 6)).astype(np.float32)
y = np.sign(x)
expect(node, inputs=[x], outputs=[y], name="test_sign")

Sin

There are 1 test cases, listed as following:

sin

node = onnx.helper.make_node(
    "Sin",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.sin(x)
expect(node, inputs=[x], outputs=[y], name="test_sin_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.sin(x)
expect(node, inputs=[x], outputs=[y], name="test_sin")

Sinh

There are 1 test cases, listed as following:

sinh

node = onnx.helper.make_node(
    "Sinh",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.sinh(x)  # expected output [-1.17520118,  0.,  1.17520118]
expect(node, inputs=[x], outputs=[y], name="test_sinh_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.sinh(x)
expect(node, inputs=[x], outputs=[y], name="test_sinh")

Size

There are 1 test cases, listed as following:

size

node = onnx.helper.make_node(
    "Size",
    inputs=["x"],
    outputs=["y"],
)

x = np.array(
    [
        [1, 2, 3],
        [4, 5, 6],
    ]
).astype(np.float32)
y = np.array(6).astype(np.int64)

expect(node, inputs=[x], outputs=[y], name="test_size_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.array(x.size).astype(np.int64)

expect(node, inputs=[x], outputs=[y], name="test_size")

Slice

There are 8 test cases, listed as following:

slice

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes", "steps"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
y = x[0:3, 0:10]
starts = np.array([0, 0], dtype=np.int64)
ends = np.array([3, 10], dtype=np.int64)
axes = np.array([0, 1], dtype=np.int64)
steps = np.array([1, 1], dtype=np.int64)

expect(
    node, inputs=[x, starts, ends, axes, steps], outputs=[y], name="test_slice"
)

slice_default_axes

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([0, 0, 3], dtype=np.int64)
ends = np.array([20, 10, 4], dtype=np.int64)
y = x[:, :, 3:4]

expect(
    node, inputs=[x, starts, ends], outputs=[y], name="test_slice_default_axes"
)

slice_default_steps

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([0, 0, 3], dtype=np.int64)
ends = np.array([20, 10, 4], dtype=np.int64)
axes = np.array([0, 1, 2], dtype=np.int64)
y = x[:, :, 3:4]

expect(
    node,
    inputs=[x, starts, ends, axes],
    outputs=[y],
    name="test_slice_default_steps",
)

slice_end_out_of_bounds

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes", "steps"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([1], dtype=np.int64)
ends = np.array([1000], dtype=np.int64)
axes = np.array([1], dtype=np.int64)
steps = np.array([1], dtype=np.int64)
y = x[:, 1:1000]

expect(
    node,
    inputs=[x, starts, ends, axes, steps],
    outputs=[y],
    name="test_slice_end_out_of_bounds",
)

slice_neg

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes", "steps"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([0], dtype=np.int64)
ends = np.array([-1], dtype=np.int64)
axes = np.array([1], dtype=np.int64)
steps = np.array([1], dtype=np.int64)
y = x[:, 0:-1]

expect(
    node,
    inputs=[x, starts, ends, axes, steps],
    outputs=[y],
    name="test_slice_neg",
)

slice_neg_steps

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes", "steps"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([20, 10, 4], dtype=np.int64)
ends = np.array([0, 0, 1], dtype=np.int64)
axes = np.array([0, 1, 2], dtype=np.int64)
steps = np.array([-1, -3, -2]).astype(np.int64)
y = x[20:0:-1, 10:0:-3, 4:1:-2]

expect(
    node,
    inputs=[x, starts, ends, axes, steps],
    outputs=[y],
    name="test_slice_neg_steps",
)

slice_negative_axes

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([0, 0, 3], dtype=np.int64)
ends = np.array([20, 10, 4], dtype=np.int64)
axes = np.array([0, -2, -1], dtype=np.int64)
y = x[:, :, 3:4]

expect(
    node,
    inputs=[x, starts, ends, axes],
    outputs=[y],
    name="test_slice_negative_axes",
)

slice_start_out_of_bounds

node = onnx.helper.make_node(
    "Slice",
    inputs=["x", "starts", "ends", "axes", "steps"],
    outputs=["y"],
)

x = np.random.randn(20, 10, 5).astype(np.float32)
starts = np.array([1000], dtype=np.int64)
ends = np.array([1000], dtype=np.int64)
axes = np.array([1], dtype=np.int64)
steps = np.array([1], dtype=np.int64)
y = x[:, 1000:1000]

expect(
    node,
    inputs=[x, starts, ends, axes, steps],
    outputs=[y],
    name="test_slice_start_out_of_bounds",
)

Softmax

There are 2 test cases, listed as following:

softmax

node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
)
x = np.array([[-1, 0, 1]]).astype(np.float32)
# expected output [[0.09003058, 0.24472848, 0.66524094]]
y = softmax(x, axis=1)
expect(node, inputs=[x], outputs=[y], name="test_softmax_example")

softmax_axis

x = np.array([[0, 1, 2, 3], [10000, 10001, 10002, 10003]]).astype(np.float32)
# expected output
# [[0.032058604 0.08714432  0.23688284  0.6439143  ]
# [0.032058604 0.08714432  0.23688284  0.6439143  ]]
y = softmax(x)

node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
)
expect(node, inputs=[x], outputs=[y], name="test_softmax_large_number")

x = np.abs(np.random.randn(3, 4, 5).astype(np.float32))
node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
    axis=0,
)
y = softmax(x, axis=0)
expect(node, inputs=[x], outputs=[y], name="test_softmax_axis_0")

node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
    axis=1,
)
y = softmax(x, axis=1)
expect(node, inputs=[x], outputs=[y], name="test_softmax_axis_1")

node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
    axis=2,
)
y = softmax(x, axis=2)
expect(node, inputs=[x], outputs=[y], name="test_softmax_axis_2")

node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
    axis=-1,
)
y = softmax(x, axis=-1)
expect(node, inputs=[x], outputs=[y], name="test_softmax_negative_axis")

# default axis is -1
node = onnx.helper.make_node(
    "Softmax",
    inputs=["x"],
    outputs=["y"],
)
expect(node, inputs=[x], outputs=[y], name="test_softmax_default_axis")

SoftmaxCrossEntropyLoss

There are 34 test cases, listed as following:

input_shape_is_NCd1_mean_weight_negative_ii

reduction = "mean"
ignore_index = np.int64(-1)

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1 = 3, 5, 6
np.random.seed(0)
x = np.random.rand(N, C, dim1).astype(np.float32)
labels = np.random.randint(0, high=C, size=(N, dim1)).astype(np.int64)
labels[0][0] = -1
weight = np.random.rand(C).astype(np.float32)

sce = softmaxcrossentropy(
    x, labels, weight=weight, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[x, labels, weight],
    outputs=[sce],
    name="test_sce_NCd1_mean_weight_negative_ii",
)

input_shape_is_NCd1_mean_weight_negative_ii_log_prob

reduction = "mean"
ignore_index = np.int64(-1)

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1 = 3, 5, 6
np.random.seed(0)
x = np.random.rand(N, C, dim1).astype(np.float32)
labels = np.random.randint(0, high=C, size=(N, dim1)).astype(np.int64)
labels[0][0] = -1
weight = np.random.rand(C).astype(np.float32)

loss, log_prob = softmaxcrossentropy(
    x,
    labels,
    weight=weight,
    reduction=reduction,
    ignore_index=ignore_index,
    get_log_prob=True,
)

expect(
    node,
    inputs=[x, labels, weight],
    outputs=[loss, log_prob],
    name="test_sce_NCd1_mean_weight_negative_ii_log_prob",
)

input_shape_is_NCd1d2d3_none_no_weight_negative_ii

reduction = "none"
ignore_index = np.int64(-5)

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1, dim2, dim3 = 3, 5, 6, 6, 5
np.random.seed(0)
x = np.random.rand(N, C, dim1, dim2, dim3).astype(np.float32)
labels = np.random.randint(0, high=C, size=(N, dim1, dim2, dim3)).astype(
    np.int64
)
labels[0][0][0][0] = -5

sce = softmaxcrossentropy(
    x, labels, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[x, labels],
    outputs=[sce],
    name="test_sce_NCd1d2d3_none_no_weight_negative_ii",
)

input_shape_is_NCd1d2d3_none_no_weight_negative_ii_log_prob

reduction = "none"
ignore_index = np.int64(-5)

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C, dim1, dim2, dim3 = 3, 5, 6, 6, 5
np.random.seed(0)
x = np.random.rand(N, C, dim1, dim2, dim3).astype(np.float32)
labels = np.random.randint(0, high=C, size=(N, dim1, dim2, dim3)).astype(
    np.int64
)
labels[0][0][0][0] = -5

loss, log_prob = softmaxcrossentropy(
    x, labels, reduction=reduction, ignore_index=ignore_index, get_log_prob=True
)

expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_NCd1d2d3_none_no_weight_negative_ii_log_prob",
)

input_shape_is_NCd1d2d3_sum_weight_high_ii

reduction = "sum"
ignore_index = np.int64(10)

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C = 3, 5
np.random.seed(0)
x = np.random.rand(N, C).astype(np.float32)
labels = np.random.randint(0, high=C, size=(N)).astype(np.int64)
labels[0] = 10
weight = np.random.rand(C).astype(np.float32)

sce = softmaxcrossentropy(
    x, labels, weight=weight, reduction=reduction, ignore_index=ignore_index
)

expect(
    node,
    inputs=[x, labels, weight],
    outputs=[sce],
    name="test_sce_NCd1d2d3_sum_weight_high_ii",
)

input_shape_is_NCd1d2d3_sum_weight_high_ii_log_prob

reduction = "sum"
ignore_index = np.int64(10)

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

N, C = 3, 5
np.random.seed(0)
x = np.random.rand(N, C).astype(np.float32)
labels = np.random.randint(0, high=C, size=(N)).astype(np.int64)
labels[0] = 10
weight = np.random.rand(C).astype(np.float32)

loss, log_prob = softmaxcrossentropy(
    x,
    labels,
    weight=weight,
    reduction=reduction,
    ignore_index=ignore_index,
    get_log_prob=True,
)

expect(
    node,
    inputs=[x, labels, weight],
    outputs=[loss, log_prob],
    name="test_sce_NCd1d2d3_sum_weight_high_ii_log_prob",
)

input_shape_is_NCd1d2d3d4d5_mean_weight

reduction = "mean"

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
)

N, C, dim1, dim2, dim3, dim4, dim5 = 3, 5, 6, 6, 5, 3, 4
np.random.seed(0)
x = np.random.rand(N, C, dim1, dim2, dim3, dim4, dim5).astype(np.float32)
labels = np.random.randint(
    0, high=C, size=(N, dim1, dim2, dim3, dim4, dim5)
).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

sce = softmaxcrossentropy(x, labels, weight=weight, reduction=reduction)

expect(
    node,
    inputs=[x, labels, weight],
    outputs=[sce],
    name="test_sce_NCd1d2d3d4d5_mean_weight",
)

input_shape_is_NCd1d2d3d4d5_mean_weight_log_prob

reduction = "mean"

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

N, C, dim1, dim2, dim3, dim4, dim5 = 3, 5, 6, 6, 5, 3, 4
np.random.seed(0)
x = np.random.rand(N, C, dim1, dim2, dim3, dim4, dim5).astype(np.float32)
labels = np.random.randint(
    0, high=C, size=(N, dim1, dim2, dim3, dim4, dim5)
).astype(np.int64)
weight = np.random.rand(C).astype(np.float32)

loss, log_prob = softmaxcrossentropy(
    x, labels, weight=weight, reduction=reduction, get_log_prob=True
)

expect(
    node,
    inputs=[x, labels, weight],
    outputs=[loss, log_prob],
    name="test_sce_NCd1d2d3d4d5_mean_weight_log_prob",
)

input_shape_is_NCd1d2d3d4d5_none_no_weight

reduction = "none"

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
)

N, C, dim1, dim2, dim3, dim4, dim5 = 3, 5, 6, 6, 5, 3, 4
np.random.seed(0)
x = np.random.rand(N, C, dim1, dim2, dim3, dim4, dim5).astype(np.float32)
labels = np.random.randint(
    0, high=C, size=(N, dim1, dim2, dim3, dim4, dim5)
).astype(np.int64)

sce = softmaxcrossentropy(x, labels, reduction=reduction)

expect(
    node,
    inputs=[x, labels],
    outputs=[sce],
    name="test_sce_NCd1d2d3d4d5_none_no_weight",
)

input_shape_is_NCd1d2d3d4d5_none_no_weight_log_prob

reduction = "none"

node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

N, C, dim1, dim2, dim3, dim4, dim5 = 3, 5, 6, 6, 5, 3, 4
np.random.seed(0)
x = np.random.rand(N, C, dim1, dim2, dim3, dim4, dim5).astype(np.float32)
labels = np.random.randint(
    0, high=C, size=(N, dim1, dim2, dim3, dim4, dim5)
).astype(np.int64)

loss, log_prob = softmaxcrossentropy(
    x, labels, reduction=reduction, get_log_prob=True
)

expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_NCd1d2d3d4d5_none_no_weight_log_prob",
)

softmaxcrossentropy_mean

# Define operator attributes.
reduction = "mean"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels)

# Check results
expect(node, inputs=[x, labels], outputs=[sce], name="test_sce_mean")

softmaxcrossentropy_mean_3d

# Define operator attributes.
reduction = "mean"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2).astype(np.float32)
y = np.random.randint(0, high=5, size=(3, 2)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, y)

# Check results
expect(node, inputs=[x, y], outputs=[sce], name="test_sce_mean_3d")

softmaxcrossentropy_mean_3d_log_prob

# Define operator attributes.
reduction = "mean"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2).astype(np.float32)
y = np.random.randint(0, high=5, size=(3, 2)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(x, y, get_log_prob=True)

# Check results
expect(
    node,
    inputs=[x, y],
    outputs=[loss, log_prob],
    name="test_sce_mean_3d_log_prob",
)

softmaxcrossentropy_mean_log_prob

# Define operator attributes.
reduction = "mean"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(x, labels, get_log_prob=True)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_mean_log_prob",
)

softmaxcrossentropy_mean_no_weights_ii

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
labels[0] = np.int64(2)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, ignore_index=ignore_index)

# Check results
expect(
    node, inputs=[x, labels], outputs=[sce], name="test_sce_mean_no_weight_ii"
)

softmaxcrossentropy_mean_no_weights_ii_3d

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2)).astype(np.int64)
labels[0][0] = np.int64(2)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, ignore_index=ignore_index)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[sce],
    name="test_sce_mean_no_weight_ii_3d",
)

softmaxcrossentropy_mean_no_weights_ii_3d_log_prob

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2)).astype(np.int64)
labels[0][0] = np.int64(2)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, ignore_index=ignore_index, get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_mean_no_weight_ii_3d_log_prob",
)

softmaxcrossentropy_mean_no_weights_ii_4d

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2, 7).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2, 7)).astype(np.int64)
labels[0][0][0] = np.int64(2)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(
    x, labels, reduction=reduction, ignore_index=ignore_index
)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[sce],
    name="test_sce_mean_no_weight_ii_4d",
)

softmaxcrossentropy_mean_no_weights_ii_4d_log_prob

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2, 7).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2, 7)).astype(np.int64)
labels[0][0][0] = np.int64(2)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, reduction=reduction, ignore_index=ignore_index, get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_mean_no_weight_ii_4d_log_prob",
)

softmaxcrossentropy_mean_no_weights_ii_log_prob

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
labels[0] = np.int64(2)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, ignore_index=ignore_index, get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_mean_no_weight_ii_log_prob",
)

softmaxcrossentropy_mean_weights

# Define operator attributes.
reduction = "mean"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
weights = np.array([0.9, 0.7, 0.8, 0.9, 0.9], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, weight=weights)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[sce],
    name="test_sce_mean_weight",
)

softmaxcrossentropy_mean_weights_ii

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(0)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
labels[0] = np.int64(0)
weights = np.array([0.9, 0.7, 0.8, 0.9, 0.9], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, weight=weights, ignore_index=ignore_index)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[sce],
    name="test_sce_mean_weight_ii",
)

softmaxcrossentropy_mean_weights_ii_3d

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(1)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2)).astype(np.int64)
labels[0][0] = np.int64(1)
weights = np.array([0.2, 0.3, 0.6, 0.1, 0.5], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, weight=weights, ignore_index=ignore_index)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[sce],
    name="test_sce_mean_weight_ii_3d",
)

softmaxcrossentropy_mean_weights_ii_3d_log_prob

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(1)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2)).astype(np.int64)
labels[0][0] = np.int64(1)
weights = np.array([0.2, 0.3, 0.6, 0.1, 0.5], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, weight=weights, ignore_index=ignore_index, get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[loss, log_prob],
    name="test_sce_mean_weight_ii_3d_log_prob",
)

softmaxcrossentropy_mean_weights_ii_4d

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2, 7).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2, 7)).astype(np.int64)
labels[0][0][0] = np.int64(2)
weights = np.array([0.2, 0.3, 0.6, 0.1, 0.5], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(
    x, labels, reduction=reduction, weight=weights, ignore_index=ignore_index
)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[sce],
    name="test_sce_mean_weight_ii_4d",
)

softmaxcrossentropy_mean_weights_ii_4d_log_prob

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(2)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5, 2, 7).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3, 2, 7)).astype(np.int64)
labels[0][0][0] = np.int64(2)
weights = np.array([0.2, 0.3, 0.6, 0.1, 0.5], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x,
    labels,
    reduction=reduction,
    weight=weights,
    ignore_index=ignore_index,
    get_log_prob=True,
)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[loss, log_prob],
    name="test_sce_mean_weight_ii_4d_log_prob",
)

softmaxcrossentropy_mean_weights_ii_log_prob

# Define operator attributes.
reduction = "mean"
ignore_index = np.int64(0)

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
    ignore_index=ignore_index,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
labels[0] = np.int64(0)
weights = np.array([0.9, 0.7, 0.8, 0.9, 0.9], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, weight=weights, ignore_index=ignore_index, get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[loss, log_prob],
    name="test_sce_mean_weight_ii_log_prob",
)

softmaxcrossentropy_mean_weights_log_prob

# Define operator attributes.
reduction = "mean"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
weights = np.array([0.9, 0.7, 0.8, 0.9, 0.9], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, weight=weights, get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[loss, log_prob],
    name="test_sce_mean_weight_log_prob",
)

softmaxcrossentropy_none

# Define operator attributes.
reduction = "none"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, reduction="none")

# Check results
expect(node, inputs=[x, labels], outputs=[sce], name="test_sce_none")

softmaxcrossentropy_none_log_prob

# Define operator attributes.
reduction = "none"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, reduction="none", get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_none_log_prob",
)

softmaxcrossentropy_none_weights

# Define operator attributes.
reduction = "none"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
weights = np.array([0.9, 0.7, 0.8, 0.9, 0.9], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, weight=weights, reduction="none")

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[sce],
    name="test_sce_none_weights",
)

softmaxcrossentropy_none_weights_log_prob

# Define operator attributes.
reduction = "none"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y", "w"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)
weights = np.array([0.9, 0.7, 0.8, 0.9, 0.9], dtype=np.float32)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, weight=weights, reduction="none", get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels, weights],
    outputs=[loss, log_prob],
    name="test_sce_none_weights_log_prob",
)

softmaxcrossentropy_sum

# Define operator attributes.
reduction = "sum"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
sce = softmaxcrossentropy(x, labels, reduction="sum")

# Check results
expect(node, inputs=[x, labels], outputs=[sce], name="test_sce_sum")

softmaxcrossentropy_sum_log_prob

# Define operator attributes.
reduction = "sum"

# Create operator.
node = onnx.helper.make_node(
    "SoftmaxCrossEntropyLoss",
    inputs=["x", "y"],
    outputs=["z", "log_prob"],
    reduction=reduction,
)

# Define operator inputs.
np.random.seed(0)
x = np.random.rand(3, 5).astype(np.float32)
labels = np.random.randint(0, high=5, size=(3,)).astype(np.int64)

# Compute SoftmaxCrossEntropyLoss
loss, log_prob = softmaxcrossentropy(
    x, labels, reduction="sum", get_log_prob=True
)

# Check results
expect(
    node,
    inputs=[x, labels],
    outputs=[loss, log_prob],
    name="test_sce_sum_log_prob",
)

Softplus

There are 1 test cases, listed as following:

softplus

node = onnx.helper.make_node(
    "Softplus",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.log(
    np.exp(x) + 1
)  # expected output [0.31326166, 0.69314718, 1.31326163]
expect(node, inputs=[x], outputs=[y], name="test_softplus_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.log(np.exp(x) + 1)
expect(node, inputs=[x], outputs=[y], name="test_softplus")

Softsign

There are 1 test cases, listed as following:

softsign

node = onnx.helper.make_node(
    "Softsign",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.array([-0.5, 0, 0.5]).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_softsign_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = x / (1 + np.abs(x))
expect(node, inputs=[x], outputs=[y], name="test_softsign")

SpaceToDepth

There are 2 test cases, listed as following:

example

node = onnx.helper.make_node(
    "SpaceToDepth",
    inputs=["x"],
    outputs=["y"],
    blocksize=2,
)

# (1, 1, 4, 6) input tensor
x = np.array(
    [
        [
            [
                [0, 6, 1, 7, 2, 8],
                [12, 18, 13, 19, 14, 20],
                [3, 9, 4, 10, 5, 11],
                [15, 21, 16, 22, 17, 23],
            ]
        ]
    ]
).astype(np.float32)

# (1, 4, 2, 3) output tensor
y = np.array(
    [
        [
            [[0, 1, 2], [3, 4, 5]],
            [[6, 7, 8], [9, 10, 11]],
            [[12, 13, 14], [15, 16, 17]],
            [[18, 19, 20], [21, 22, 23]],
        ]
    ]
).astype(np.float32)
expect(node, inputs=[x], outputs=[y], name="test_spacetodepth_example")

spacetodepth

b, c, h, w = shape = (2, 2, 6, 6)
blocksize = 2
node = onnx.helper.make_node(
    "SpaceToDepth",
    inputs=["x"],
    outputs=["y"],
    blocksize=blocksize,
)
x = np.random.random_sample(shape).astype(np.float32)
tmp = np.reshape(
    x, [b, c, h // blocksize, blocksize, w // blocksize, blocksize]
)
tmp = np.transpose(tmp, [0, 3, 5, 1, 2, 4])
y = np.reshape(tmp, [b, c * (blocksize**2), h // blocksize, w // blocksize])
expect(node, inputs=[x], outputs=[y], name="test_spacetodepth")

Split

There are 10 test cases, listed as following:

1d_opset13

node_input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)

node = onnx.helper.make_node(
    "Split",
    inputs=["input"],
    outputs=["output_1", "output_2", "output_3"],
    axis=0,
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0]).astype(np.float32),
    np.array([5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_equal_parts_1d_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

split = np.array([2, 4]).astype(np.int64)
node = onnx.helper.make_node(
    "Split",
    inputs=["input", "split"],
    outputs=["output_1", "output_2"],
    axis=0,
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_variable_parts_1d_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

1d_opset18

node_input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)

node = onnx.helper.make_node(
    "Split",
    inputs=["input"],
    outputs=["output_1", "output_2", "output_3"],
    axis=0,
    num_outputs=3,
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0]).astype(np.float32),
    np.array([5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_equal_parts_1d_opset18",
)

split = np.array([2, 4]).astype(np.int64)
node = onnx.helper.make_node(
    "Split",
    inputs=["input", "split"],
    outputs=["output_1", "output_2"],
    axis=0,
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_variable_parts_1d_opset18",
)

1d_uneven_split_opset18

node_input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]).astype(np.float32)

# If axis is not specified, split is applied on default axis 0
node = onnx.helper.make_node(
    "Split",
    inputs=["input"],
    outputs=["output_1", "output_2", "output_3", "output_4"],
    num_outputs=4,
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0]).astype(np.float32),
    np.array([5.0, 6.0]).astype(np.float32),
    np.array([7.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_1d_uneven_split_opset18",
)

2d_opset13

node_input = np.array(
    [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [7.0, 8.0, 9.0, 10.0, 11.0, 12.0]]
).astype(np.float32)

node = onnx.helper.make_node(
    "Split", inputs=["input"], outputs=["output_1", "output_2"], axis=1
)

expected_outputs = [
    np.array([[1.0, 2.0, 3.0], [7.0, 8.0, 9.0]]).astype(np.float32),
    np.array([[4.0, 5.0, 6.0], [10.0, 11.0, 12.0]]).astype(np.float32),
]

expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_equal_parts_2d_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

split = np.array([2, 4]).astype(np.int64)
node = onnx.helper.make_node(
    "Split",
    inputs=["input", "split"],
    outputs=["output_1", "output_2"],
    axis=1,
)

expected_outputs = [
    np.array([[1.0, 2.0], [7.0, 8.0]]).astype(np.float32),
    np.array([[3.0, 4.0, 5.0, 6.0], [9.0, 10.0, 11.0, 12.0]]).astype(
        np.float32
    ),
]

expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_variable_parts_2d_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

2d_opset18

node_input = np.array(
    [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [7.0, 8.0, 9.0, 10.0, 11.0, 12.0]]
).astype(np.float32)

node = onnx.helper.make_node(
    "Split",
    inputs=["input"],
    outputs=["output_1", "output_2"],
    axis=1,
    num_outputs=2,
)

expected_outputs = [
    np.array([[1.0, 2.0, 3.0], [7.0, 8.0, 9.0]]).astype(np.float32),
    np.array([[4.0, 5.0, 6.0], [10.0, 11.0, 12.0]]).astype(np.float32),
]

expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_equal_parts_2d",
)

split = np.array([2, 4]).astype(np.int64)
node = onnx.helper.make_node(
    "Split",
    inputs=["input", "split"],
    outputs=["output_1", "output_2"],
    axis=1,
)

expected_outputs = [
    np.array([[1.0, 2.0], [7.0, 8.0]]).astype(np.float32),
    np.array([[3.0, 4.0, 5.0, 6.0], [9.0, 10.0, 11.0, 12.0]]).astype(
        np.float32
    ),
]

expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_variable_parts_2d_opset18",
)

2d_uneven_split_opset18

node_input = np.array(
    [
        [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
        [9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0],
    ]
).astype(np.float32)

node = onnx.helper.make_node(
    "Split",
    inputs=["input"],
    outputs=["output_1", "output_2", "output_3"],
    axis=1,
    num_outputs=3,
)

expected_outputs = [
    np.array([[1.0, 2.0, 3.0], [9.0, 10.0, 11.0]]).astype(np.float32),
    np.array([[4.0, 5.0, 6.0], [12.0, 13.0, 14.0]]).astype(np.float32),
    np.array([[7.0, 8.0], [15.0, 16.0]]).astype(np.float32),
]

expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_2d_uneven_split_opset18",
)

default_values_opset13

node_input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)

# If axis is not specified, split is applied on default axis 0
node = onnx.helper.make_node(
    "Split", inputs=["input"], outputs=["output_1", "output_2", "output_3"]
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0]).astype(np.float32),
    np.array([5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_equal_parts_default_axis_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

split = np.array([2, 4]).astype(np.int64)
node = onnx.helper.make_node(
    "Split", inputs=["input", "split"], outputs=["output_1", "output_2"]
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_variable_parts_default_axis_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

default_values_opset18

node_input = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).astype(np.float32)

# If axis is not specified, split is applied on default axis 0
node = onnx.helper.make_node(
    "Split",
    inputs=["input"],
    outputs=["output_1", "output_2", "output_3"],
    num_outputs=3,
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0]).astype(np.float32),
    np.array([5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input],
    outputs=expected_outputs,
    name="test_split_equal_parts_default_axis_opset18",
)

split = np.array([2, 4]).astype(np.int64)
node = onnx.helper.make_node(
    "Split", inputs=["input", "split"], outputs=["output_1", "output_2"]
)

expected_outputs = [
    np.array([1.0, 2.0]).astype(np.float32),
    np.array([3.0, 4.0, 5.0, 6.0]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_variable_parts_default_axis_opset18",
)

zero_size_splits_opset13

# 1-dimensional tensor with dimension_size=0
node_input = np.array([]).astype(np.float32)

# Split empty tensor to tensors of size zero
split = np.array([0, 0, 0]).astype(np.int64)
node = onnx.helper.make_node(
    "Split",
    inputs=["input", "split"],
    outputs=["output_1", "output_2", "output_3"],
)

expected_outputs = [
    np.array([]).astype(np.float32),
    np.array([]).astype(np.float32),
    np.array([]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_zero_size_splits_opset13",
    opset_imports=[onnx.helper.make_opsetid("", 13)],
)

zero_size_splits_opset18

# 1-dimensional tensor with dimension_size=0
node_input = np.array([]).astype(np.float32)

# Split empty tensor to tensors of size zero
split = np.array([0, 0, 0]).astype(np.int64)
node = onnx.helper.make_node(
    "Split",
    inputs=["input", "split"],
    outputs=["output_1", "output_2", "output_3"],
)

expected_outputs = [
    np.array([]).astype(np.float32),
    np.array([]).astype(np.float32),
    np.array([]).astype(np.float32),
]
expect(
    node,
    inputs=[node_input, split],
    outputs=expected_outputs,
    name="test_split_zero_size_splits_opset18",
)

SplitToSequence

There are 3 test cases, listed as following:

nokeepdims

data = np.arange(18).reshape((3, 6)).astype(np.float32)

node = onnx.helper.make_node(
    "SplitToSequence",
    ["data"],
    ["seq"],
    axis=1,
    keepdims=0,
)

expected_outputs = [[data[:, i] for i in range(data.shape[1])]]

expect(
    node,
    inputs=[data],
    outputs=expected_outputs,
    name="test_split_to_sequence_nokeepdims",
)

with_split_1

data = np.arange(18).reshape((3, 6)).astype(np.float32)
split = np.array(2, dtype=np.int64)

node = onnx.helper.make_node(
    "SplitToSequence", ["data", "split"], ["seq"], axis=1
)

expected_outputs = [
    [
        np.array([[0.0, 1.0], [6.0, 7.0], [12.0, 13.0]], dtype=np.float32),
        np.array([[2.0, 3.0], [8.0, 9.0], [14.0, 15.0]], dtype=np.float32),
        np.array([[4.0, 5.0], [10.0, 11.0], [16.0, 17.0]], dtype=np.float32),
    ]
]

expect(
    node,
    inputs=[data, split],
    outputs=expected_outputs,
    name="test_split_to_sequence_1",
)

with_split_2

data = np.arange(18).reshape((3, 6)).astype(np.float32)
split = np.array([1, 2], dtype=np.int64)

node = onnx.helper.make_node(
    "SplitToSequence", ["data", "split"], ["seq"], axis=0
)

expected_outputs = [
    [
        data[:1],
        data[1:],
    ]
]

expect(
    node,
    inputs=[data, split],
    outputs=expected_outputs,
    name="test_split_to_sequence_2",
)

Sqrt

There are 1 test cases, listed as following:

sqrt

node = onnx.helper.make_node(
    "Sqrt",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([1, 4, 9]).astype(np.float32)
y = np.sqrt(x)  # expected output [1., 2., 3.]
expect(node, inputs=[x], outputs=[y], name="test_sqrt_example")

x = np.abs(np.random.randn(3, 4, 5).astype(np.float32))
y = np.sqrt(x)
expect(node, inputs=[x], outputs=[y], name="test_sqrt")

Squeeze

There are 2 test cases, listed as following:

squeeze

node = onnx.helper.make_node(
    "Squeeze",
    inputs=["x", "axes"],
    outputs=["y"],
)
x = np.random.randn(1, 3, 4, 5).astype(np.float32)
axes = np.array([0], dtype=np.int64)
y = np.squeeze(x, axis=0)

expect(node, inputs=[x, axes], outputs=[y], name="test_squeeze")

squeeze_negative_axes

node = onnx.helper.make_node(
    "Squeeze",
    inputs=["x", "axes"],
    outputs=["y"],
)
x = np.random.randn(1, 3, 1, 5).astype(np.float32)
axes = np.array([-2], dtype=np.int64)
y = np.squeeze(x, axis=-2)
expect(node, inputs=[x, axes], outputs=[y], name="test_squeeze_negative_axes")

StringConcat

There are 1 test cases, listed as following:

stringconcat

node = onnx.helper.make_node(
    "StringConcat",
    inputs=["x", "y"],
    outputs=["result"],
)
x = np.array(["abc", "def"]).astype("object")
y = np.array([".com", ".net"]).astype("object")
result = np.array(["abc.com", "def.net"]).astype("object")

expect(node, inputs=[x, y], outputs=[result], name="test_string_concat")

x = np.array(["cat", "dog", "snake"]).astype("object")
y = np.array(["s"]).astype("object")
result = np.array(["cats", "dogs", "snakes"]).astype("object")

expect(
    node,
    inputs=[x, y],
    outputs=[result],
    name="test_string_concat_broadcasting",
)

x = np.array("cat").astype("object")
y = np.array("s").astype("object")
result = np.array("cats").astype("object")

expect(
    node,
    inputs=[x, y],
    outputs=[result],
    name="test_string_concat_zero_dimensional",
)

x = np.array(["abc", ""]).astype("object")
y = np.array(["", "abc"]).astype("object")
result = np.array(["abc", "abc"]).astype("object")

expect(
    node,
    inputs=[x, y],
    outputs=[result],
    name="test_string_concat_empty_string",
)

x = np.array(["的", "中"]).astype("object")
y = np.array(["的", "中"]).astype("object")
result = np.array(["的的", "中中"]).astype("object")

expect(
    node,
    inputs=[x, y],
    outputs=[result],
    name="test_string_concat_utf8",
)

StringNormalizer

There are 6 test cases, listed as following:

monday_casesensintive_lower

input = np.array(["monday", "tuesday", "wednesday", "thursday"]).astype(object)
output = np.array(["tuesday", "wednesday", "thursday"]).astype(object)
stopwords = ["monday"]

node = onnx.helper.make_node(
    "StringNormalizer",
    inputs=["x"],
    outputs=["y"],
    case_change_action="LOWER",
    is_case_sensitive=1,
    stopwords=stopwords,
)
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_strnormalizer_export_monday_casesensintive_lower",
)

monday_casesensintive_nochangecase

input = np.array(["monday", "tuesday", "wednesday", "thursday"]).astype(object)
output = np.array(["tuesday", "wednesday", "thursday"]).astype(object)
stopwords = ["monday"]

node = onnx.helper.make_node(
    "StringNormalizer",
    inputs=["x"],
    outputs=["y"],
    is_case_sensitive=1,
    stopwords=stopwords,
)
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_strnormalizer_export_monday_casesensintive_nochangecase",
)

monday_casesensintive_upper

input = np.array(["monday", "tuesday", "wednesday", "thursday"]).astype(object)
output = np.array(["TUESDAY", "WEDNESDAY", "THURSDAY"]).astype(object)
stopwords = ["monday"]

node = onnx.helper.make_node(
    "StringNormalizer",
    inputs=["x"],
    outputs=["y"],
    case_change_action="UPPER",
    is_case_sensitive=1,
    stopwords=stopwords,
)
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_strnormalizer_export_monday_casesensintive_upper",
)

monday_empty_output

input = np.array(["monday", "monday"]).astype(object)
output = np.array([""]).astype(object)
stopwords = ["monday"]

node = onnx.helper.make_node(
    "StringNormalizer",
    inputs=["x"],
    outputs=["y"],
    case_change_action="UPPER",
    is_case_sensitive=1,
    stopwords=stopwords,
)
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_strnormalizer_export_monday_empty_output",
)

monday_insensintive_upper_twodim

input = (
    np.array(
        ["Monday", "tuesday", "wednesday", "Monday", "tuesday", "wednesday"]
    )
    .astype(object)
    .reshape([1, 6])
)

# It does upper case cecedille, accented E
# and german umlaut but fails
# with german eszett
output = (
    np.array(["TUESDAY", "WEDNESDAY", "TUESDAY", "WEDNESDAY"])
    .astype(object)
    .reshape([1, 4])
)
stopwords = ["monday"]

node = onnx.helper.make_node(
    "StringNormalizer",
    inputs=["x"],
    outputs=["y"],
    case_change_action="UPPER",
    stopwords=stopwords,
)
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_strnormalizer_export_monday_insensintive_upper_twodim",
)

nostopwords_nochangecase

input = np.array(["monday", "tuesday"]).astype(object)
output = input

# No stopwords. This is a NOOP
node = onnx.helper.make_node(
    "StringNormalizer",
    inputs=["x"],
    outputs=["y"],
    is_case_sensitive=1,
)
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_strnormalizer_nostopwords_nochangecase",
)

StringSplit

There are 5 test cases, listed as following:

basic

node = onnx.helper.make_node(
    "StringSplit",
    inputs=["x"],
    outputs=["substrings", "length"],
    delimiter=".",
    maxsplit=None,
)

x = np.array(["abc.com", "def.net"]).astype(object)

substrings = np.array([["abc", "com"], ["def", "net"]]).astype(object)

length = np.array([2, 2], dtype=np.int64)

expect(
    node,
    inputs=[x],
    outputs=[substrings, length],
    name="test_string_split_basic",
)

consecutive_delimiters

node = onnx.helper.make_node(
    "StringSplit",
    inputs=["x"],
    outputs=["substrings", "length"],
    delimiter="-",
    maxsplit=None,
)

x = np.array(["o-n-n--x-", "o-n----nx"]).astype(object)

substrings = np.array(
    [["o", "n", "n", "", "x", ""], ["o", "n", "", "", "", "nx"]]
).astype(object)

length = np.array([6, 6], dtype=np.int64)

expect(
    node,
    inputs=[x],
    outputs=[substrings, length],
    name="test_string_split_consecutive_delimiters",
)

empty_string_delimiter

for delimiter, test_name in (
    ("", "test_string_split_empty_string_delimiter"),
    (None, "test_string_split_no_delimiter"),
):
    node = onnx.helper.make_node(
        "StringSplit",
        inputs=["x"],
        outputs=["substrings", "length"],
        delimiter=delimiter,
        maxsplit=None,
    )

    x = np.array(
        ["hello world !", "  hello   world !", " hello world   ! "]
    ).astype(object)

    substrings = np.array(
        [
            ["hello", "world", "!"],
            ["hello", "world", "!"],
            ["hello", "world", "!"],
        ]
    ).astype(object)

    length = np.array([3, 3, 3], dtype=np.int64)

    expect(
        node,
        inputs=[x],
        outputs=[substrings, length],
        name=test_name,
    )

empty_string_split

node = onnx.helper.make_node(
    "StringSplit",
    inputs=["x"],
    outputs=["substrings", "length"],
    delimiter=None,
    maxsplit=None,
)

x = np.array([]).astype(object)

substrings = np.array([]).astype(object).reshape(0, 0)

length = np.array([], dtype=np.int64)

expect(
    node,
    inputs=[x],
    outputs=[substrings, length],
    name="test_string_split_empty_tensor",
    output_type_protos=[
        onnx.helper.make_tensor_type_proto(onnx.TensorProto.STRING, (0, None)),
        None,
    ],
)

maxsplit

node = onnx.helper.make_node(
    "StringSplit",
    inputs=["x"],
    outputs=["substrings", "length"],
    maxsplit=2,
)

x = np.array(
    [["hello world", "def.net"], ["o n n x", "the quick brown fox"]]
).astype(object)

substrings = np.array(
    [
        [["hello", "world", ""], ["def.net", "", ""]],
        [["o", "n", "n x"], ["the", "quick", "brown fox"]],
    ]
).astype(object)

length = np.array([[2, 1], [3, 3]], np.int64)

expect(
    node,
    inputs=[x],
    outputs=[substrings, length],
    name="test_string_split_maxsplit",
)

Sub

There are 2 test cases, listed as following:

sub

node = onnx.helper.make_node(
    "Sub",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.array([1, 2, 3]).astype(np.float32)
y = np.array([3, 2, 1]).astype(np.float32)
z = x - y  # expected output [-2., 0., 2.]
expect(node, inputs=[x, y], outputs=[z], name="test_sub_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(3, 4, 5).astype(np.float32)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub")

x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.int8)
y = np.random.randint(12, size=(3, 4, 5), dtype=np.int8)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_int8")

x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.int16)
y = np.random.randint(12, size=(3, 4, 5), dtype=np.int16)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_int16")

x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint8)
y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint8)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint8")

x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint16)
y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint16)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint16")

x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint32)
y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint32)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint32")

x = np.random.randint(12, 24, size=(3, 4, 5), dtype=np.uint64)
y = np.random.randint(12, size=(3, 4, 5), dtype=np.uint64)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_uint64")

sub_broadcast

node = onnx.helper.make_node(
    "Sub",
    inputs=["x", "y"],
    outputs=["z"],
)

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.random.randn(5).astype(np.float32)
z = x - y
expect(node, inputs=[x, y], outputs=[z], name="test_sub_bcast")

Sum

There are 1 test cases, listed as following:

sum

data_0 = np.array([3, 0, 2]).astype(np.float32)
data_1 = np.array([1, 3, 4]).astype(np.float32)
data_2 = np.array([2, 6, 6]).astype(np.float32)
result = np.array([6, 9, 12]).astype(np.float32)
node = onnx.helper.make_node(
    "Sum",
    inputs=["data_0", "data_1", "data_2"],
    outputs=["result"],
)
expect(
    node,
    inputs=[data_0, data_1, data_2],
    outputs=[result],
    name="test_sum_example",
)

node = onnx.helper.make_node(
    "Sum",
    inputs=["data_0"],
    outputs=["result"],
)
expect(node, inputs=[data_0], outputs=[data_0], name="test_sum_one_input")

result = np.add(data_0, data_1)
node = onnx.helper.make_node(
    "Sum",
    inputs=["data_0", "data_1"],
    outputs=["result"],
)
expect(
    node, inputs=[data_0, data_1], outputs=[result], name="test_sum_two_inputs"
)

Swish

There are 1 test cases, listed as following:

swish

node = onnx.helper.make_node(
    "Swish",
    inputs=["x"],
    outputs=["y"],
    alpha=1.0,  # pass alpha as attribute
)

x = np.array([3, 4, 5], dtype=np.float32)
y = swish(x, alpha=1.0)

expect(
    node,
    inputs=[x],
    outputs=[y],
    name="test_swish",
    opset_imports=[onnx.helper.make_opsetid("", 24)],
)

Tan

There are 1 test cases, listed as following:

tan

node = onnx.helper.make_node(
    "Tan",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.tan(x)
expect(node, inputs=[x], outputs=[y], name="test_tan_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.tan(x)
expect(node, inputs=[x], outputs=[y], name="test_tan")

Tanh

There are 1 test cases, listed as following:

tanh

node = onnx.helper.make_node(
    "Tanh",
    inputs=["x"],
    outputs=["y"],
)

x = np.array([-1, 0, 1]).astype(np.float32)
y = np.tanh(x)  # expected output [-0.76159418, 0., 0.76159418]
expect(node, inputs=[x], outputs=[y], name="test_tanh_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.tanh(x)
expect(node, inputs=[x], outputs=[y], name="test_tanh")

TensorScatter

There are 3 test cases, listed as following:

tensorscatter

node = onnx.helper.make_node(
    "TensorScatter",
    inputs=["past_cache", "update", "write_indices"],
    outputs=["present_cache"],
    mode="linear",
)
past_cache = np.array(
    [
        [[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
        [[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
    ],
    dtype=np.float32,
)
update = np.array(
    [
        [[[5, 5, 5, 5, 5]]],
        [[[1, 1, 1, 1, 1]]],
    ],
    dtype=np.float32,
)
write_indices = np.array([1, 2], dtype=np.int64)
present_cache = np.array(
    [
        [[[1, 2, 3, 4, 5], [5, 5, 5, 5, 5], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
        [[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [1, 1, 1, 1, 1], [4, 3, 2, 1, 0]]],
    ],
    dtype=np.float32,
)
expect(
    node,
    inputs=[past_cache, update, write_indices],
    outputs=[present_cache],
    name="test_tensorscatter",
)

tensorscatter_3d

node = onnx.helper.make_node(
    "TensorScatter",
    inputs=["past_cache", "update", "write_indices"],
    outputs=["present_cache"],
)
past_cache = np.array(
    [
        [
            [1, 2, 3, 4, 5],
            [5, 6, 7, 8, 9],
            [8, 7, 6, 5, 4],
            [5, 4, 3, 2, 1],
        ],
        [
            [1, 2, 3, 4, 5],
            [5, 6, 7, 8, 9],
            [8, 7, 6, 5, 4],
            [5, 4, 3, 2, 1],
        ],
        [
            [1, 2, 3, 4, 5],
            [5, 6, 7, 8, 9],
            [8, 7, 6, 5, 4],
            [5, 4, 3, 2, 1],
        ],
    ],
    dtype=np.float32,
)
update = np.array(
    [
        [
            [4, 4, 4, 4, 4],
            [5, 5, 5, 5, 5],
        ],
        [
            [6, 6, 6, 6, 6],
            [7, 7, 7, 7, 7],
        ],
        [
            [2, 2, 2, 2, 2],
            [3, 3, 3, 3, 3],
        ],
    ],
    dtype=np.float32,
)
write_indices = np.array([1, 2, 0], dtype=np.int64)
present_cache = np.array(
    [
        [
            [1, 2, 3, 4, 5],
            [4, 4, 4, 4, 4],
            [5, 5, 5, 5, 5],
            [5, 4, 3, 2, 1],
        ],
        [
            [1, 2, 3, 4, 5],
            [5, 6, 7, 8, 9],
            [6, 6, 6, 6, 6],
            [7, 7, 7, 7, 7],
        ],
        [
            [2, 2, 2, 2, 2],
            [3, 3, 3, 3, 3],
            [8, 7, 6, 5, 4],
            [5, 4, 3, 2, 1],
        ],
    ],
    dtype=np.float32,
)
expect(
    node,
    inputs=[past_cache, update, write_indices],
    outputs=[present_cache],
    name="test_tensorscatter_3d",
)

tensorscatter_circular

node = onnx.helper.make_node(
    "TensorScatter",
    inputs=["past_cache", "update", "write_indices"],
    outputs=["present_cache"],
    mode="circular",
)
past_cache = np.array(
    [
        [[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
        [[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
    ],
    dtype=np.float32,
)
update = np.array(
    [
        [
            [
                [5, 5, 5, 5, 5],
                [6, 6, 6, 6, 6],
            ]
        ],
        [
            [
                [1, 1, 1, 1, 1],
                [2, 2, 2, 2, 2],
            ]
        ],
    ],
    dtype=np.float32,
)
write_indices = np.array([1, 3], dtype=np.int64)
present_cache = np.array(
    [
        [[[1, 2, 3, 4, 5], [5, 5, 5, 5, 5], [6, 6, 6, 6, 6], [4, 3, 2, 1, 0]]],
        [[[2, 2, 2, 2, 2], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [1, 1, 1, 1, 1]]],
    ],
    dtype=np.float32,
)
expect(
    node,
    inputs=[past_cache, update, write_indices],
    outputs=[present_cache],
    name="test_tensorscatter_circular",
)

TfIdfVectorizer

There are 7 test cases, listed as following:

tf_batch_onlybigrams_skip0

input = np.array([[1, 1, 3, 3, 3, 7], [8, 6, 7, 5, 6, 8]]).astype(np.int32)
output = np.array(
    [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0]]
).astype(np.float32)

ngram_counts = np.array([0, 4]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=2,
    max_gram_length=2,
    max_skip_count=0,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_batch_onlybigrams_skip0",
)

tf_batch_onlybigrams_skip5

input = np.array([[1, 1, 3, 3, 3, 7], [8, 6, 7, 5, 6, 8]]).astype(np.int32)
output = np.array(
    [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]]
).astype(np.float32)

ngram_counts = np.array([0, 4]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=2,
    max_gram_length=2,
    max_skip_count=5,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_batch_onlybigrams_skip5",
)

tf_batch_uniandbigrams_skip5

input = np.array([[1, 1, 3, 3, 3, 7], [8, 6, 7, 5, 6, 8]]).astype(np.int32)
output = np.array(
    [[0.0, 3.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0]]
).astype(np.float32)

ngram_counts = np.array([0, 4]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=1,
    max_gram_length=2,
    max_skip_count=5,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_batch_uniandbigrams_skip5",
)

tf_only_bigrams_skip0

input = np.array([1, 1, 3, 3, 3, 7, 8, 6, 7, 5, 6, 8]).astype(np.int32)
output = np.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]).astype(np.float32)

ngram_counts = np.array([0, 4]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=2,
    max_gram_length=2,
    max_skip_count=0,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_only_bigrams_skip0",
)

tf_onlybigrams_levelempty

input = np.array([1, 1, 3, 3, 3, 7, 8, 6, 7, 5, 6, 8]).astype(np.int32)
output = np.array([1.0, 1.0, 1.0]).astype(np.float32)

ngram_counts = np.array([0, 0]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2]).astype(np.int64)
pool_int64s = np.array([5, 6, 7, 8, 6, 7]).astype(  # unigrams none
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=2,
    max_gram_length=2,
    max_skip_count=0,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_onlybigrams_levelempty",
)

tf_onlybigrams_skip5

input = np.array([1, 1, 3, 3, 3, 7, 8, 6, 7, 5, 6, 8]).astype(np.int32)
output = np.array([0.0, 0.0, 0.0, 0.0, 1.0, 3.0, 1.0]).astype(np.float32)

ngram_counts = np.array([0, 4]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=2,
    max_gram_length=2,
    max_skip_count=5,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_onlybigrams_skip5",
)

tf_uniandbigrams_skip5

input = np.array([1, 1, 3, 3, 3, 7, 8, 6, 7, 5, 6, 8]).astype(np.int32)
output = np.array([0.0, 3.0, 1.0, 0.0, 1.0, 3.0, 1.0]).astype(np.float32)

ngram_counts = np.array([0, 4]).astype(np.int64)
ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
    np.int64
)  # bigrams

helper = TfIdfVectorizerHelper(
    mode="TF",
    min_gram_length=1,
    max_gram_length=2,
    max_skip_count=5,
    ngram_counts=ngram_counts,
    ngram_indexes=ngram_indexes,
    pool_int64s=pool_int64s,
)
node = helper.make_node_noweights()
expect(
    node,
    inputs=[input],
    outputs=[output],
    name="test_tfidfvectorizer_tf_uniandbigrams_skip5",
)

ThresholdedRelu

There are 2 test cases, listed as following:

default

default_alpha = 1.0
node = onnx.helper.make_node("ThresholdedRelu", inputs=["x"], outputs=["y"])
x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, default_alpha, np.inf)
y[y == default_alpha] = 0

expect(node, inputs=[x], outputs=[y], name="test_thresholdedrelu_default")

thresholdedrelu

alpha = 2.0
node = onnx.helper.make_node(
    "ThresholdedRelu", inputs=["x"], outputs=["y"], alpha=alpha
)

x = np.array([-1.5, 0.0, 1.2, 2.0, 2.2]).astype(np.float32)
y = np.clip(x, alpha, np.inf)  # expected output [0., 0., 0., 0., 2.2]
y[y == alpha] = 0

expect(node, inputs=[x], outputs=[y], name="test_thresholdedrelu_example")

x = np.random.randn(3, 4, 5).astype(np.float32)
y = np.clip(x, alpha, np.inf)
y[y == alpha] = 0

expect(node, inputs=[x], outputs=[y], name="test_thresholdedrelu")

Tile

There are 2 test cases, listed as following:

tile

node = onnx.helper.make_node("Tile", inputs=["x", "y"], outputs=["z"])

x = np.random.rand(2, 3, 4, 5).astype(np.float32)

repeats = np.random.randint(low=1, high=10, size=(np.ndim(x),)).astype(np.int64)

z = np.tile(x, repeats)

expect(node, inputs=[x, repeats], outputs=[z], name="test_tile")

tile_precomputed

node = onnx.helper.make_node("Tile", inputs=["x", "y"], outputs=["z"])

x = np.array([[0, 1], [2, 3]], dtype=np.float32)

repeats = np.array([2, 2], dtype=np.int64)

z = np.array(
    [[0, 1, 0, 1], [2, 3, 2, 3], [0, 1, 0, 1], [2, 3, 2, 3]], dtype=np.float32
)

expect(node, inputs=[x, repeats], outputs=[z], name="test_tile_precomputed")

TopK

There are 7 test cases, listed as following:

top_k

axis = 1
largest = 1

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [
        [0, 1, 2, 3],
        [4, 5, 6, 7],
        [8, 9, 10, 11],
    ],
    dtype=np.float32,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 3.  2.  1.]
# [ 7.  6.  5.]
# [11. 10.  9.]]
# print(indices_ref)
# [[3 2 1]
# [3 2 1]
# [3 2 1]]

expect(
    node, inputs=[X, K], outputs=[values_ref, indices_ref], name="test_top_k"
)

top_k_negative_axis

axis = -1
largest = 1

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [
        [0, 1, 2, 3],
        [4, 5, 6, 7],
        [8, 9, 10, 11],
    ],
    dtype=np.float32,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 3.  2.  1.]
# [ 7.  6.  5.]
# [11. 10.  9.]]
# print(indices_ref)
# [[3 2 1]
# [3 2 1]
# [3 2 1]]

expect(
    node,
    inputs=[X, K],
    outputs=[values_ref, indices_ref],
    name="test_top_k_negative_axis",
)

top_k_same_values

axis = 0
largest = 0

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [0, 0, 0, 0],
    dtype=np.int64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# (Pdb) print(values_ref)
# [0 0 0]
# (Pdb) print(indices_ref)
# [0 1 2]

expect(
    node,
    inputs=[X, K],
    outputs=[values_ref, indices_ref],
    name="test_top_k_same_values",
)

top_k_same_values_2d

axis = 1
largest = 1

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [[0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 1, 1]],
    dtype=np.int64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[0 0 0]
# [1 1 1]
# [1 1 2]]
# print(indices_ref)
# [[0 1 2]
# [0 1 2]
# [2 3 0]]

expect(
    node,
    inputs=[X, K],
    outputs=[values_ref, indices_ref],
    name="test_top_k_same_values_2d",
)

top_k_same_values_largest

axis = 0
largest = 1

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [0, 0, 0, 0],
    dtype=np.int64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [0 0 0]
# print(indices_ref)
# [0 1 2]

expect(
    node,
    inputs=[X, K],
    outputs=[values_ref, indices_ref],
    name="test_top_k_same_values_largest",
)

top_k_smallest

axis = 1
largest = 0
sorted_ = 1
k = 3

node = onnx.helper.make_node(
    "TopK",
    inputs=["x", "k"],
    outputs=["values", "indices"],
    axis=axis,
    largest=largest,
    sorted=sorted_,
)

X = np.array(
    [
        [0, 1, 2, 3],
        [4, 5, 6, 7],
        [11, 10, 9, 8],
    ],
    dtype=np.float32,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 0.  1.  2.]
# [ 4.  5.  6.]
# [ 8.  9. 10.]]
# print(indices_ref)
# [[0 1 2]
# [0 1 2]
# [3 2 1]]

expect(
    node,
    inputs=[X, K],
    outputs=[values_ref, indices_ref],
    name="test_top_k_smallest",
)

top_k_uint64

axis = 1
largest = 1

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [
        [0, 1, 2, 3],
        [4, 5, 6, 7],
        [8, 9, 10, 11],
    ],
    dtype=np.uint64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 3  2  1]
# [ 7  6  5]
# [11 10  9]]
# print(indices_ref)
# [[3 2 1]
# [3 2 1]
# [3 2 1]]

expect(
    node,
    inputs=[X, K],
    outputs=[values_ref, indices_ref],
    name="test_top_k_uint64",
)

Transpose

There are 2 test cases, listed as following:

all_permutations

shape = (2, 3, 4)
data = np.random.random_sample(shape).astype(np.float32)
permutations = list(itertools.permutations(np.arange(len(shape))))

for i, permutation in enumerate(permutations):
    node = onnx.helper.make_node(
        "Transpose",
        inputs=["data"],
        outputs=["transposed"],
        perm=permutation,
    )
    transposed = np.transpose(data, permutation)
    expect(
        node,
        inputs=[data],
        outputs=[transposed],
        name=f"test_transpose_all_permutations_{i}",
    )

default

shape = (2, 3, 4)
data = np.random.random_sample(shape).astype(np.float32)

node = onnx.helper.make_node(
    "Transpose", inputs=["data"], outputs=["transposed"]
)

transposed = np.transpose(data)
expect(node, inputs=[data], outputs=[transposed], name="test_transpose_default")

Trilu

There are 18 test cases, listed as following:

tril

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[4, 0, 0, 0, 0],
#   [1, 2, 0, 0, 0],
#   [9, 4, 1, 0, 0],
#   [4, 3, 4, 2, 0]]
y = tril_reference_implementation(x)
expect(node, inputs=[x], outputs=[y], name="test_tril")

tril_neg

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(-1).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[0, 0, 0, 0, 0],
#   [1, 0, 0, 0, 0],
#   [9, 4, 0, 0, 0],
#   [4, 3, 4, 0, 0]]
y = tril_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_tril_neg")

tril_one_row

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(3, 1, 5)).astype(np.int64)
# X:
# [[[6, 2, 4, 1, 6]],
#
#  [[8, 3, 8, 7, 0]],
#
#  [[2, 2, 9, 5, 9]]]
# expect result:
# [[[6, 0, 0, 0, 0]],
#
#  [[8, 0, 0, 0, 0]],
#
#  [[2, 0, 0, 0, 0]]]
y = tril_reference_implementation(x)
expect(node, inputs=[x], outputs=[y], name="test_tril_one_row_neg")

tril_out_neg

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(-7).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[0, 0, 0, 0, 0],
#   [0, 0, 0, 0, 0],
#   [0, 0, 0, 0, 0],
#   [0, 0, 0, 0, 0]]
y = tril_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_tril_out_neg")

tril_out_pos

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
    upper=0,
)
x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(6).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
y = tril_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_tril_out_pos")

tril_pos

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(2).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[4, 7, 3, 0, 0],
#   [1, 2, 8, 6, 0],
#   [9, 4, 1, 8, 7],
#   [4, 3, 4, 2, 4]]
y = tril_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_tril_pos")

tril_square

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(2, 3, 3)).astype(np.int64)
# X:
# [[[0, 4, 3],
#   [2, 0, 9],
#   [8, 2, 5]],
#
#  [[2, 7, 2],
#   [2, 6, 0],
#   [2, 6, 5]]]
# expect result:
# [[[0, 0, 0],
#   [2, 0, 0],
#   [8, 2, 5]],
#
#  [[2, 0, 0],
#   [2, 6, 0],
#   [2, 6, 5]]]
y = tril_reference_implementation(x)
expect(node, inputs=[x], outputs=[y], name="test_tril_square")

tril_square_neg

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(2, 3, 3)).astype(np.int64)
k = np.array(-1).astype(np.int64)
# X:
# [[[0, 4, 3],
#   [2, 0, 9],
#   [8, 2, 5]],
#
#  [[2, 7, 2],
#   [2, 6, 0],
#   [2, 6, 5]]]
# expect result:
# [[[0, 0, 0],
#   [2, 0, 0],
#   [8, 2, 0]],
#
#  [[0, 0, 0],
#   [2, 0, 0],
#   [2, 6, 0]]]
y = tril_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_tril_square_neg")

tril_zero

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
    upper=0,
)

x = np.random.randint(10, size=(3, 0, 5)).astype(np.int64)
k = np.array(6).astype(np.int64)
# X:
# []
# expect result:
# []
y = tril_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_tril_zero")

triu

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x"],
    outputs=["y"],
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 0, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[4, 7, 3, 7, 9],
#   [0, 2, 8, 6, 9],
#   [0, 0, 0, 8, 7],
#   [0, 0, 0, 2, 4]]
y = triu_reference_implementation(x)
expect(node, inputs=[x], outputs=[y], name="test_triu")

triu_neg

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(-1).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 0, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [0, 4, 0, 8, 7],
#   [0, 0, 4, 2, 4]]
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_neg")

triu_one_row

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(3, 1, 5)).astype(np.int64)
k = np.array(1).astype(np.int64)
# X:
# [[[1, 4, 9, 7, 1]],
#
#  [[9, 2, 8, 8, 4]],
#
#  [[3, 9, 7, 4, 2]]]
# expect result:
# [[[0, 4, 9, 7, 1]],
#
#  [[0, 2, 8, 8, 4]],
#
#  [[0, 9, 7, 4, 2]]]
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_one_row")

triu_out_neg_out

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(-7).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 0, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 0, 8, 7],
#   [4, 3, 4, 2, 4]]
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_out_neg_out")

triu_out_pos

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(6).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 0, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[0, 0, 0, 0, 0],
#   [0, 0, 0, 0, 0],
#   [0, 0, 0, 0, 0],
#   [0, 0, 0, 0, 0]]
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_out_pos")

triu_pos

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(4, 5)).astype(np.int64)
k = np.array(2).astype(np.int64)
# X:
#  [[4, 7, 3, 7, 9],
#   [1, 2, 8, 6, 9],
#   [9, 4, 0, 8, 7],
#   [4, 3, 4, 2, 4]]
# expect result:
#  [[0, 0, 3, 7, 9],
#   [0, 0, 0, 6, 9],
#   [0, 0, 0, 0, 7],
#   [0, 0, 0, 0, 0]]
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_pos")

triu_square

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x"],
    outputs=["y"],
)

x = np.random.randint(10, size=(2, 3, 3)).astype(np.int64)
y = triu_reference_implementation(x)
# X:
# [[[4, 6, 9],
#   [7, 5, 4],
#   [8, 1, 2]],
#
#  [[1, 4, 9],
#   [9, 6, 3],
#   [8, 9, 8]]]
# expect result:
# [[[4, 6, 9],
#   [0, 5, 4],
#   [0, 0, 2]],
#
#  [[1, 4, 9],
#   [0, 6, 3],
#   [0, 0, 8]]]
expect(node, inputs=[x], outputs=[y], name="test_triu_square")

triu_square_neg

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(2, 3, 3)).astype(np.int64)
k = np.array(-1).astype(np.int64)
# X:
# [[[4, 6, 9],
#   [7, 5, 4],
#   [8, 1, 2]],
#
#  [[1, 4, 9],
#   [9, 6, 3],
#   [8, 9, 8]]]
# expect result:
# [[[4, 6, 9],
#   [7, 5, 4],
#   [0, 1, 2]],
#
#  [[1, 4, 9],
#   [9, 6, 3],
#   [0, 9, 8]]]
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_square_neg")

triu_zero

node = onnx.helper.make_node(
    "Trilu",
    inputs=["x", "k"],
    outputs=["y"],
)

x = np.random.randint(10, size=(0, 5)).astype(np.int64)
k = np.array(6).astype(np.int64)
# X:
# []
# expect result:
# []
y = triu_reference_implementation(x, int(k))
expect(node, inputs=[x, k], outputs=[y], name="test_triu_zero")

Unique

There are 6 test cases, listed as following:

length_1

node_sorted = onnx.helper.make_node(
    "Unique",
    inputs=["X"],
    outputs=["Y", "indices", "inverse_indices", "counts"],
    sorted=1,
)

x = np.array([0], dtype=np.int64)
y, indices, inverse_indices, counts = np.unique(x, True, True, True)
indices, inverse_indices, counts = specify_int64(
    indices, inverse_indices, counts
)
# behavior changed with numpy >= 2.0
inverse_indices = inverse_indices.reshape(-1)
# print(y)
# [0]
# print(indices)
# [0]
# print(inverse_indices)
# [0]
# print(counts)
# [1]

expect(
    node_sorted,
    inputs=[x],
    outputs=[y, indices, inverse_indices, counts],
    name="test_unique_length_1",
)

not_sorted_without_axis

node_not_sorted = onnx.helper.make_node(
    "Unique",
    inputs=["X"],
    outputs=["Y", "indices", "inverse_indices", "counts"],
    sorted=0,
)
# numpy unique does not retain original order (it sorts the output unique values)
# https://github.com/numpy/numpy/issues/8621
# we need to recover unsorted output and indices
x = np.array([2.0, 1.0, 1.0, 3.0, 4.0, 3.0], dtype=np.float32)
y, indices, inverse_indices, counts = np.unique(x, True, True, True)

# prepare index mapping from sorted to unsorted
argsorted_indices = np.argsort(indices)
inverse_indices_map = dict(
    zip(argsorted_indices, np.arange(len(argsorted_indices)), strict=True)
)

indices = indices[argsorted_indices]
y = np.take(x, indices, axis=0)
inverse_indices = np.asarray(
    [inverse_indices_map[i] for i in inverse_indices], dtype=np.int64
)
counts = counts[argsorted_indices]
indices, inverse_indices, counts = specify_int64(
    indices, inverse_indices, counts
)
# print(y)
# [2.0, 1.0, 3.0, 4.0]
# print(indices)
# [0 1 3 4]
# print(inverse_indices)
# [0, 1, 1, 2, 3, 2]
# print(counts)
# [1, 2, 2, 1]

expect(
    node_not_sorted,
    inputs=[x],
    outputs=[y, indices, inverse_indices, counts],
    name="test_unique_not_sorted_without_axis",
)

sorted_with_axis

node_sorted = onnx.helper.make_node(
    "Unique",
    inputs=["X"],
    outputs=["Y", "indices", "inverse_indices", "counts"],
    sorted=1,
    axis=0,
)

x = np.array([[1, 0, 0], [1, 0, 0], [2, 3, 4]], dtype=np.float32)
y, indices, inverse_indices, counts = np.unique(x, True, True, True, axis=0)
indices, inverse_indices, counts = specify_int64(
    indices, inverse_indices, counts
)
# behavior changed with numpy >= 2.0
inverse_indices = inverse_indices.reshape(-1)
# print(y)
# [[1. 0. 0.]
#  [2. 3. 4.]]
# print(indices)
# [0 2]
# print(inverse_indices)
# [0 0 1]
# print(counts)
# [2 1]

expect(
    node_sorted,
    inputs=[x],
    outputs=[y, indices, inverse_indices, counts],
    name="test_unique_sorted_with_axis",
)

sorted_with_axis_3d

node_sorted = onnx.helper.make_node(
    "Unique",
    inputs=["X"],
    outputs=["Y", "indices", "inverse_indices", "counts"],
    sorted=1,
    axis=1,
)

x = np.array(
    [
        [[1.0, 1.0], [0.0, 1.0], [2.0, 1.0], [0.0, 1.0]],
        [[1.0, 1.0], [0.0, 1.0], [2.0, 1.0], [0.0, 1.0]],
    ],
    dtype=np.float32,
)
y, indices, inverse_indices, counts = np.unique(x, True, True, True, axis=1)
indices, inverse_indices, counts = specify_int64(
    indices, inverse_indices, counts
)
# behavior changed with numpy >= 2.0
inverse_indices = inverse_indices.reshape(-1)
# print(y)
# [[[0. 1.]
#  [1. 1.]
#  [2. 1.]]
# [[0. 1.]
#  [1. 1.]
#  [2. 1.]]]
# print(indices)
# [1 0 2]
# print(inverse_indices)
# [1 0 2 0]
# print(counts)
# [2 1 1]
expect(
    node_sorted,
    inputs=[x],
    outputs=[y, indices, inverse_indices, counts],
    name="test_unique_sorted_with_axis_3d",
)

sorted_with_negative_axis

node_sorted = onnx.helper.make_node(
    "Unique",
    inputs=["X"],
    outputs=["Y", "indices", "inverse_indices", "counts"],
    sorted=1,
    axis=-1,
)

x = np.array([[1, 0, 0], [1, 0, 0], [2, 3, 3]], dtype=np.float32)
y, indices, inverse_indices, counts = np.unique(x, True, True, True, axis=-1)
indices, inverse_indices, counts = specify_int64(
    indices, inverse_indices, counts
)
# behavior changed with numpy >= 2.0
inverse_indices = inverse_indices.reshape(-1)
# print(y)
# [[0. 1.]
#  [0. 1.]
#  [3. 2.]]
# print(indices)
# [1 0]
# print(inverse_indices)
# [1 0 0]
# print(counts)
# [2 1]

expect(
    node_sorted,
    inputs=[x],
    outputs=[y, indices, inverse_indices, counts],
    name="test_unique_sorted_with_negative_axis",
)

sorted_without_axis

node_sorted = onnx.helper.make_node(
    "Unique",
    inputs=["X"],
    outputs=["Y", "indices", "inverse_indices", "counts"],
)

x = np.array([2.0, 1.0, 1.0, 3.0, 4.0, 3.0], dtype=np.float32)
y, indices, inverse_indices, counts = np.unique(x, True, True, True)
indices, inverse_indices, counts = specify_int64(
    indices, inverse_indices, counts
)
expect(
    node_sorted,
    inputs=[x],
    outputs=[y, indices, inverse_indices, counts],
    name="test_unique_sorted_without_axis",
)

Unsqueeze

There are 5 test cases, listed as following:

unsqueeze_negative_axes

node = onnx.helper.make_node(
    "Unsqueeze",
    inputs=["x", "axes"],
    outputs=["y"],
)
x = np.random.randn(1, 3, 1, 5).astype(np.float32)
axes = np.array([-2]).astype(np.int64)
y = np.expand_dims(x, axis=-2)
expect(node, inputs=[x, axes], outputs=[y], name="test_unsqueeze_negative_axes")

unsqueeze_one_axis

x = np.random.randn(3, 4, 5).astype(np.float32)

for i in range(x.ndim):
    axes = np.array([i]).astype(np.int64)
    node = onnx.helper.make_node(
        "Unsqueeze",
        inputs=["x", "axes"],
        outputs=["y"],
    )
    y = np.expand_dims(x, axis=i)

    expect(
        node,
        inputs=[x, axes],
        outputs=[y],
        name="test_unsqueeze_axis_" + str(i),
    )

unsqueeze_three_axes

x = np.random.randn(3, 4, 5).astype(np.float32)
axes = np.array([2, 4, 5]).astype(np.int64)

node = onnx.helper.make_node(
    "Unsqueeze",
    inputs=["x", "axes"],
    outputs=["y"],
)
y = np.expand_dims(x, axis=2)
y = np.expand_dims(y, axis=4)
y = np.expand_dims(y, axis=5)

expect(node, inputs=[x, axes], outputs=[y], name="test_unsqueeze_three_axes")

unsqueeze_two_axes

x = np.random.randn(3, 4, 5).astype(np.float32)
axes = np.array([1, 4]).astype(np.int64)

node = onnx.helper.make_node(
    "Unsqueeze",
    inputs=["x", "axes"],
    outputs=["y"],
)
y = np.expand_dims(x, axis=1)
y = np.expand_dims(y, axis=4)

expect(node, inputs=[x, axes], outputs=[y], name="test_unsqueeze_two_axes")

unsqueeze_unsorted_axes

x = np.random.randn(3, 4, 5).astype(np.float32)
axes = np.array([5, 4, 2]).astype(np.int64)

node = onnx.helper.make_node(
    "Unsqueeze",
    inputs=["x", "axes"],
    outputs=["y"],
)
y = np.expand_dims(x, axis=2)
y = np.expand_dims(y, axis=4)
y = np.expand_dims(y, axis=5)

expect(node, inputs=[x, axes], outputs=[y], name="test_unsqueeze_unsorted_axes")

Upsample

There are 1 test cases, listed as following:

nearest

node = onnx.helper.make_node(
    "Upsample",
    inputs=["X", "scales"],
    outputs=["Y"],
    mode="nearest",
)

data = np.array(
    [
        [
            [
                [1, 2],
                [3, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

scales = np.array([1.0, 1.0, 2.0, 3.0], dtype=np.float32)

output = np.array(
    [
        [
            [
                [1, 1, 1, 2, 2, 2],
                [1, 1, 1, 2, 2, 2],
                [3, 3, 3, 4, 4, 4],
                [3, 3, 3, 4, 4, 4],
            ]
        ]
    ],
    dtype=np.float32,
)

expect(
    node,
    inputs=[data, scales],
    outputs=[output],
    name="test_upsample_nearest",
    opset_imports=[helper.make_opsetid("", 9)],
)

Where

There are 2 test cases, listed as following:

long

node = onnx.helper.make_node(
    "Where",
    inputs=["condition", "x", "y"],
    outputs=["z"],
)

condition = np.array([[1, 0], [1, 1]], dtype=bool)
x = np.array([[1, 2], [3, 4]], dtype=np.int64)
y = np.array([[9, 8], [7, 6]], dtype=np.int64)
z = np.where(condition, x, y)  # expected output [[1, 8], [3, 4]]
expect(
    node, inputs=[condition, x, y], outputs=[z], name="test_where_long_example"
)

where

node = onnx.helper.make_node(
    "Where",
    inputs=["condition", "x", "y"],
    outputs=["z"],
)

condition = np.array([[1, 0], [1, 1]], dtype=bool)
x = np.array([[1, 2], [3, 4]], dtype=np.float32)
y = np.array([[9, 8], [7, 6]], dtype=np.float32)
z = np.where(condition, x, y)  # expected output [[1, 8], [3, 4]]
expect(node, inputs=[condition, x, y], outputs=[z], name="test_where_example")

Xor

There are 2 test cases, listed as following:

xor

node = onnx.helper.make_node(
    "Xor",
    inputs=["x", "y"],
    outputs=["xor"],
)

# 2d
x = (np.random.randn(3, 4) > 0).astype(bool)
y = (np.random.randn(3, 4) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor2d")

# 3d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(3, 4, 5) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor3d")

# 4d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor4d")

xor_broadcast

node = onnx.helper.make_node(
    "Xor",
    inputs=["x", "y"],
    outputs=["xor"],
)

# 3d vs 1d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(5) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor_bcast3v1d")

# 3d vs 2d
x = (np.random.randn(3, 4, 5) > 0).astype(bool)
y = (np.random.randn(4, 5) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor_bcast3v2d")

# 4d vs 2d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(5, 6) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor_bcast4v2d")

# 4d vs 3d
x = (np.random.randn(3, 4, 5, 6) > 0).astype(bool)
y = (np.random.randn(4, 5, 6) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor_bcast4v3d")

# 4d vs 4d
x = (np.random.randn(1, 4, 1, 6) > 0).astype(bool)
y = (np.random.randn(3, 1, 5, 6) > 0).astype(bool)
z = np.logical_xor(x, y)
expect(node, inputs=[x, y], outputs=[z], name="test_xor_bcast4v4d")

💔No Cover Common Operators

ConcatFromSequence (call for test cases)

GlobalLpPool (call for test cases)

GreaterOrEqual (call for test cases)

LessOrEqual (call for test cases)

MaxRoiPool (call for test cases)

Multinomial (random generator operator)

Optional (call for test cases)

OptionalGetElement (call for test cases)

RandomNormal (random generator operator)

RandomNormalLike (random generator operator)

RandomUniform (random generator operator)

RandomUniformLike (random generator operator)

SequenceAt (call for test cases)

SequenceConstruct (call for test cases)

SequenceEmpty (call for test cases)

SequenceErase (call for test cases)

SequenceLength (call for test cases)

💚Covered Experimental Operators

FlexAttention

There are 11 test cases, listed as following:

flexattention

"""Basic FlexAttention test with default settings."""
node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)

B, Hq, L, E = 2, 4, 8, 16
S, Ev = 6, 16

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

(Y,) = _compute_flex_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_causal_mask

"""FlexAttention with causal masking score_mod (Qwen-3, Gemma-3, Llama-3 pattern)."""
score_mod_graph = _make_score_mod_causal_mask_graph(TensorProto.FLOAT)

node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)
score_mod_attr = helper.make_attribute("score_mod", score_mod_graph)
node.attribute.append(score_mod_attr)

B, Hq, L, E = 1, 2, 4, 8
S, Ev = 4, 8

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

# Manually compute expected output with causal masking
scale = 1.0 / np.sqrt(E)
scores = np.einsum("bhle,bhse->bhls", Q, K) * scale
# Apply causal mask: set future positions to -inf
q_idx = np.arange(L).reshape(1, 1, L, 1)
k_idx = np.arange(S).reshape(1, 1, 1, S)
mask = q_idx >= k_idx
scores = np.where(mask, scores, -np.inf)
probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
probs = probs / probs.sum(axis=-1, keepdims=True)
Y = np.einsum("bhls,bhsv->bhlv", probs, V).astype(np.float32)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_causal_mask",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_diff_head_sizes

"""FlexAttention with different head sizes for Q/K vs V."""
node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)

B, Hq, L, E = 2, 4, 8, 16
S, Ev = 6, 32  # V has different head size

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

(Y,) = _compute_flex_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_diff_head_sizes",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_double

"""FlexAttention with double precision inputs."""
node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)

B, Hq, L, E = 2, 4, 8, 16
S, Ev = 6, 16

Q = np.random.rand(B, Hq, L, E).astype(np.float64)
K = np.random.rand(B, Hq, S, E).astype(np.float64)
V = np.random.rand(B, Hq, S, Ev).astype(np.float64)

(Y,) = _compute_flex_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_double",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_fp16

"""FlexAttention with float16 inputs."""
node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)

B, Hq, L, E = 2, 4, 8, 16
S, Ev = 6, 16

Q = np.random.rand(B, Hq, L, E).astype(np.float16)
K = np.random.rand(B, Hq, S, E).astype(np.float16)
V = np.random.rand(B, Hq, S, Ev).astype(np.float16)

(Y,) = _compute_flex_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_fp16",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_gqa

"""FlexAttention with Grouped Query Attention (GQA)."""
node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)

B, Hq, Hkv, L, S, E, Ev = 2, 8, 2, 4, 6, 16, 16

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hkv, S, E).astype(np.float32)
V = np.random.rand(B, Hkv, S, Ev).astype(np.float32)

(Y,) = _compute_flex_attention(Q, K, V)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_gqa",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_prob_mod

"""FlexAttention with prob_mod subgraph (scales probabilities)."""
scale_value = 0.5
prob_mod_graph = _make_prob_mod_scale_graph(scale_value, TensorProto.FLOAT)

node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)
prob_mod_attr = helper.make_attribute("prob_mod", prob_mod_graph)
node.attribute.append(prob_mod_attr)

B, Hq, L, E = 1, 2, 3, 4
S, Ev = 3, 4

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

scale = 1.0 / np.sqrt(E)
scores = np.einsum("bhle,bhse->bhls", Q, K) * scale
probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
probs = probs / probs.sum(axis=-1, keepdims=True)
probs = probs * scale_value
Y = np.einsum("bhls,bhsv->bhlv", probs, V).astype(np.float32)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_prob_mod",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_relative_positional

"""FlexAttention with relative positional bias score_mod."""
score_mod_graph = _make_score_mod_relative_positional_graph(TensorProto.FLOAT)

node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)
score_mod_attr = helper.make_attribute("score_mod", score_mod_graph)
node.attribute.append(score_mod_attr)

B, Hq, L, E = 1, 2, 4, 8
S, Ev = 4, 8

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

# Manually compute expected output with relative positional bias
scale = 1.0 / np.sqrt(E)
scores = np.einsum("bhle,bhse->bhls", Q, K) * scale
q_idx = np.arange(L).reshape(-1, 1)
k_idx = np.arange(S).reshape(1, -1)
rel_pos = (q_idx - k_idx).astype(np.float32)
scores = scores + rel_pos
probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
probs = probs / probs.sum(axis=-1, keepdims=True)
Y = np.einsum("bhls,bhsv->bhlv", probs, V).astype(np.float32)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_relative_positional",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_scaled

"""FlexAttention with explicit scale attribute."""
scale = 0.1
node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    scale=scale,
    domain=AI_ONNX_PREVIEW_DOMAIN,
)

B, Hq, L, E = 2, 4, 8, 16
S, Ev = 6, 16

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

(Y,) = _compute_flex_attention(Q, K, V, scale=scale)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_scaled",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_score_mod

"""FlexAttention with score_mod subgraph (adds bias to scores)."""
bias_value = 0.5
score_mod_graph = _make_score_mod_bias_graph(bias_value, TensorProto.FLOAT)

node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)
# Add score_mod as a graph attribute
score_mod_attr = helper.make_attribute("score_mod", score_mod_graph)
node.attribute.append(score_mod_attr)

B, Hq, L, E = 1, 2, 3, 4
S, Ev = 3, 4

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

scale = 1.0 / np.sqrt(E)
scores = np.einsum("bhle,bhse->bhls", Q, K) * scale
scores = scores + bias_value  # score_mod: add bias
probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
probs = probs / probs.sum(axis=-1, keepdims=True)
Y = np.einsum("bhls,bhsv->bhlv", probs, V).astype(np.float32)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_score_mod",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

flexattention_soft_cap

"""FlexAttention with soft capping score_mod (Gemma-2 pattern)."""
cap_value = 20.0
score_mod_graph = _make_score_mod_soft_cap_graph(cap_value, TensorProto.FLOAT)

node = helper.make_node(
    "FlexAttention",
    inputs=["Q", "K", "V"],
    outputs=["Y"],
    domain=AI_ONNX_PREVIEW_DOMAIN,
)
score_mod_attr = helper.make_attribute("score_mod", score_mod_graph)
node.attribute.append(score_mod_attr)

B, Hq, L, E = 1, 2, 4, 8
S, Ev = 4, 8

Q = np.random.rand(B, Hq, L, E).astype(np.float32)
K = np.random.rand(B, Hq, S, E).astype(np.float32)
V = np.random.rand(B, Hq, S, Ev).astype(np.float32)

# Manually compute expected output with soft capping
scale = 1.0 / np.sqrt(E)
scores = np.einsum("bhle,bhse->bhls", Q, K) * scale
scores = np.tanh(scores / cap_value) * cap_value
probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
probs = probs / probs.sum(axis=-1, keepdims=True)
Y = np.einsum("bhls,bhsv->bhlv", probs, V).astype(np.float32)

expect(
    node,
    inputs=[Q, K, V],
    outputs=[Y],
    name="test_flexattention_soft_cap",
    opset_imports=[
        helper.make_opsetid("", 26),
        helper.make_opsetid(AI_ONNX_PREVIEW_DOMAIN, 1),
    ],
)

💔No Cover Experimental Operators

Model Test Coverage

bvlc_alexnet

bvlc_alexnet has 40 nodes. Of these, 40 are covered by node tests (100.0%)

nodes

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 1 kernel_shape: 3 pads: 3 strides: 2

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 2 storage_order: 0 strides: 1

densenet121

densenet121 has 1746 nodes. Of these, 1746 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 1 pads: 1 strides: 1

BatchNormalization: 1 out of 3 attributes covered

epsilon: 1 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 1 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 3 storage_order: 0 strides: 1

Unsqueeze: 1 out of 0 attributes covered

inception_v1

inception_v1 has 237 nodes. Of these, 237 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 2 pads: 2 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 1 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 1 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 3 storage_order: 0 strides: 2

Unsqueeze: 1 out of 0 attributes covered

inception_v2

inception_v2 has 916 nodes. Of these, 916 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 3 pads: 3 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 1 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 1 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 3 storage_order: 0 strides: 2

Unsqueeze: 1 out of 0 attributes covered

resnet50

resnet50 has 415 nodes. Of these, 415 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 3 pads: 3 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 2 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 1 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 3 storage_order: 0 strides: 2

Unsqueeze: 1 out of 0 attributes covered

shufflenet

shufflenet has 446 nodes. Of these, 446 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 3 pads: 3 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 2 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 6 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 3 storage_order: 0 strides: 2

Transpose: 1 out of 1 attributes covered

perm: 1

Unsqueeze: 1 out of 0 attributes covered

squeezenet_old

squeezenet_old has 105 nodes. Of these, 105 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 3 pads: 3 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 2 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 6 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 1 pads: 3 storage_order: 0 strides: 2

Transpose: 1 out of 1 attributes covered

perm: 1

Unsqueeze: 1 out of 0 attributes covered

vgg19

vgg19 has 82 nodes. Of these, 82 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 3 pads: 3 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 2 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 6 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 1 beta: 1 bias: 1 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 2 pads: 3 storage_order: 0 strides: 2

Transpose: 1 out of 1 attributes covered

perm: 1

Unsqueeze: 1 out of 0 attributes covered

zfnet512

zfnet512 has 38 nodes. Of these, 38 are covered by node tests (100.0%)

nodes

AveragePool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 count_include_pad: 0 dilations: 0 kernel_shape: 3 pads: 3 strides: 2

BatchNormalization: 1 out of 3 attributes covered

epsilon: 2 momentum: 0 training_mode: 0

Concat: 1 out of 1 attributes covered

axis: 1

ConstantOfShape: 1 out of 1 attributes covered

value: 1

Conv: 4 out of 6 attributes covered

auto_pad: 0 dilations: 0 group: 6 kernel_shape: 5 pads: 4 strides: 3

Dropout: 1 out of 1 attributes covered

seed: 0

Gemm: 1 out of 4 attributes covered

alpha: 0 beta: 0 transA: 0 transB: 1

LRN: 4 out of 4 attributes covered

alpha: 2 beta: 1 bias: 2 size: 1

MaxPool: 3 out of 7 attributes covered

auto_pad: 0 ceil_mode: 0 dilations: 0 kernel_shape: 2 pads: 3 storage_order: 0 strides: 2

Transpose: 1 out of 1 attributes covered

perm: 1

Unsqueeze: 1 out of 0 attributes covered

Overall Test Coverage

To be filled.