Skip to content

Commit 540bee7

Browse files
vacu9708Youngsik Yang
authored andcommitted
[Relax][ONNX] Update ReduceL1 to version 18
- Update ReduceL1-13 to ReduceL1-18 - Add the corresponding test cases
1 parent 3c7c515 commit 540bee7

File tree

2 files changed

+162
-14
lines changed

2 files changed

+162
-14
lines changed

python/tvm/relax/frontend/onnx/onnx_frontend.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2612,6 +2612,29 @@ def _impl_v13(cls, bb, inputs, attr, params):
26122612
keepdims = attr.get("keepdims", 1)
26132613
return relax.op.sum(relax.op.abs(data), axes, keepdims)
26142614

2615+
@classmethod
2616+
def _impl_v18(cls, bb, inputs, attr, params):
2617+
data = inputs[0]
2618+
keepdims = attr.get("keepdims", 1)
2619+
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
2620+
2621+
# Optional axes input
2622+
axes = None
2623+
if len(inputs) > 1 and inputs[1] is not None:
2624+
axes_const = get_constant(inputs[1], params)
2625+
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
2626+
axes = axes_const.data.numpy().tolist()
2627+
2628+
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
2629+
if not axes and not noop_with_empty_axes:
2630+
return relax.op.sum(relax.op.abs(data), None, keepdims)
2631+
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
2632+
elif not axes and noop_with_empty_axes:
2633+
return data
2634+
# If axes is provided, reduce over specified axes
2635+
else:
2636+
return relax.op.sum(relax.op.abs(data), axes, keepdims)
2637+
26152638

26162639
class ReduceL2(OnnxOpConverter):
26172640
"""Converts an onnx ReduceL2 node into an equivalent Relax expression."""

tests/python/relax/test_frontend_onnx.py

Lines changed: 139 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1503,24 +1503,24 @@ def verify_embedlayernormalization(
15031503
)
15041504

15051505

1506-
def create_reduce_test_parameters():
1506+
def create_reduce_test_parameters_axes_attr():
15071507
output = []
15081508
for value in [True, False]:
1509-
output.append(("ReduceMax", value))
1510-
output.append(("ReduceMean", value))
1511-
output.append(("ReduceMin", value))
1512-
output.append(("ReduceProd", value))
1513-
output.append(("ReduceSum", value))
1514-
output.append(("ReduceSumSquare", value))
1515-
output.append(("ReduceLogSum", value))
1516-
output.append(("ReduceLogSumExp", value))
1517-
output.append(("ReduceL1", value))
1518-
output.append(("ReduceL2", value))
1509+
output.append(("ReduceMax", value, 11))
1510+
output.append(("ReduceMean", value, 13))
1511+
output.append(("ReduceMin", value, 11))
1512+
output.append(("ReduceProd", value, 13))
1513+
output.append(("ReduceSum", value, 11))
1514+
output.append(("ReduceSumSquare", value, 13))
1515+
output.append(("ReduceLogSum", value, 13))
1516+
output.append(("ReduceLogSumExp", value, 13))
1517+
output.append(("ReduceL1", value, 13))
1518+
output.append(("ReduceL2", value, 13))
15191519
return output
15201520

15211521

1522-
@pytest.mark.parametrize("func, dynamic", create_reduce_test_parameters())
1523-
def test_all_reduce_funcs(func, dynamic):
1522+
@pytest.mark.parametrize("func, dynamic, opset", create_reduce_test_parameters_axes_attr())
1523+
def test_all_reduce_funcs_axes_attr(func, dynamic, opset):
15241524
def verify_reduce_func(func, data, axis, keepdims):
15251525
inshape = data.shape
15261526
outshape = np.sum(data, axis=axis, keepdims=keepdims == 1).shape
@@ -1549,7 +1549,7 @@ def verify_reduce_func(func, data, axis, keepdims):
15491549

15501550
inputs_dict = {"x": data}
15511551
# Reduction ops accumulate arithmetic errors, so we use a higher tolerance.
1552-
check_correctness(model, inputs_dict, opset=11, rtol=1e-4, atol=1e-4)
1552+
check_correctness(model, inputs_dict, opset=opset, rtol=1e-4, atol=1e-4)
15531553

15541554
for keepdims in [True, False]:
15551555
verify_reduce_func(
@@ -1577,6 +1577,131 @@ def verify_reduce_func(func, data, axis, keepdims):
15771577
)
15781578

15791579

1580+
def create_reduce_test_parameters_axes_input():
1581+
output = []
1582+
for dynamic in [True, False]:
1583+
# TODO(@vacu9708): Enable the tests after implementing other reduce ops
1584+
# output.append(("ReduceMax", dynamic, 20))
1585+
# output.append(("ReduceMean", dynamic, 18))
1586+
# output.append(("ReduceMin", dynamic, 20))
1587+
# output.append(("ReduceProd", dynamic, 18))
1588+
# output.append(("ReduceSum", dynamic, 13))
1589+
# output.append(("ReduceSumSquare", dynamic, 18))
1590+
# output.append(("ReduceLogSum", dynamic, 18))
1591+
# output.append(("ReduceLogSumExp", dynamic, 18))
1592+
output.append(("ReduceL1", dynamic, 18))
1593+
# output.append(("ReduceL2", dynamic, 18))
1594+
return output
1595+
1596+
1597+
@pytest.mark.parametrize("func, dynamic, opset", create_reduce_test_parameters_axes_input())
1598+
def test_all_reduce_funcs_axes_input(func, dynamic, opset):
1599+
def verify_reduce_func(func, data, axes, keepdims, noop_with_empty_axes):
1600+
inshape = data.shape
1601+
1602+
inputs = ["x"]
1603+
initializers = []
1604+
1605+
# Optional `axes` input
1606+
if axes is not None:
1607+
axes_name = "reduce_axes"
1608+
axes_np = np.asarray(axes, dtype=np.int64)
1609+
axes_init = helper.make_tensor(
1610+
name=axes_name,
1611+
data_type=TensorProto.INT64,
1612+
dims=axes_np.shape,
1613+
vals=axes_np,
1614+
)
1615+
initializers.append(axes_init)
1616+
inputs.append(axes_name)
1617+
1618+
# Determine input and output shapes
1619+
if not axes and not noop_with_empty_axes:
1620+
outshape = np.sum(data, axis=None, keepdims=keepdims).shape
1621+
elif not axes and noop_with_empty_axes:
1622+
outshape = inshape
1623+
else:
1624+
outshape = np.sum(data, axis=axes, keepdims=keepdims).shape
1625+
1626+
if dynamic:
1627+
in_list = ["?"] * len(inshape)
1628+
out_list = ["?"] * len(outshape)
1629+
else:
1630+
in_list = list(inshape)
1631+
out_list = list(outshape)
1632+
1633+
# Make a model node
1634+
node = helper.make_node(
1635+
func,
1636+
inputs=inputs,
1637+
outputs=["y"],
1638+
keepdims=keepdims,
1639+
noop_with_empty_axes=noop_with_empty_axes,
1640+
)
1641+
1642+
# Make a model graph and a model
1643+
graph = helper.make_graph(
1644+
[node],
1645+
"reduce18_test",
1646+
inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, in_list)],
1647+
initializer=initializers,
1648+
outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_list)],
1649+
)
1650+
model = helper.make_model(graph, producer_name="reduce18_test")
1651+
1652+
# Run TVM importer vs onnxruntime
1653+
inputs_dict = {"x": data}
1654+
check_correctness(model, inputs_dict, opset=opset, rtol=1e-4, atol=1e-4)
1655+
1656+
# Verify
1657+
for keepdims in [True, False]:
1658+
# no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all dimensions.
1659+
verify_reduce_func(
1660+
func,
1661+
np.random.randn(3, 2, 2).astype(np.float32),
1662+
axes=[],
1663+
keepdims=keepdims,
1664+
noop_with_empty_axes=False,
1665+
)
1666+
1667+
# no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all dimensions.
1668+
verify_reduce_func(
1669+
func,
1670+
np.random.randn(3, 2, 2).astype(np.float32),
1671+
axes=None,
1672+
keepdims=keepdims,
1673+
noop_with_empty_axes=False,
1674+
)
1675+
1676+
# no `axes` input && `noop_with_empty_axes` = 1 -> return the input unchanged.
1677+
verify_reduce_func(
1678+
func,
1679+
np.random.randn(4, 3).astype(np.float32),
1680+
axes=[],
1681+
keepdims=keepdims,
1682+
noop_with_empty_axes=True,
1683+
)
1684+
1685+
# no `axes` input && `noop_with_empty_axes` = 1 -> return the input unchanged.
1686+
# (onnxruntime bug) Runtime error on the onnxruntime part
1687+
# verify_reduce_func(
1688+
# func,
1689+
# np.random.randn(4, 3).astype(np.float32),
1690+
# axes=None,
1691+
# keepdims=keepdims,
1692+
# noop_with_empty_axes=True,
1693+
# )
1694+
1695+
# `axes` provided -> reduce over specified axes.
1696+
verify_reduce_func(
1697+
func,
1698+
np.random.randn(3, 3, 3, 1).astype(np.float32),
1699+
axes=(1, 2),
1700+
keepdims=keepdims,
1701+
noop_with_empty_axes=True,
1702+
)
1703+
1704+
15801705
@pytest.mark.parametrize("in_dtype", [np.float32, np.int32])
15811706
@pytest.mark.parametrize("axis", [None, 0, 1, 2])
15821707
@pytest.mark.parametrize("keepdims", [None, True, False])

0 commit comments

Comments
 (0)