diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 7e8616f65f7b..8a5eec070412 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -728,10 +728,10 @@ class Mod(BinaryBase): @classmethod def _impl_v10(cls, bb, inputs, attr, params): if attr.get("fmod", 0) == 0: - cls.numpy_op = _np.fmod + cls.numpy_op = _np.mod cls.relax_op = relax.op.floor_mod else: - cls.numpy_op = _np.mod + cls.numpy_op = _np.fmod cls.relax_op = relax.op.mod return cls.base_impl(bb, inputs, attr, params) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 09f3a7b01d88..c7a1d4f2390c 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -362,7 +362,7 @@ def verify_binary_scalar(op_name, attrs={}, domain=None, dtype=TensorProto.INT32 "Mul": np.multiply, "Div": np.divide, "Pow": np.power, - "Mod": np.mod if attrs.get("fmod", 0) else np.fmod, + "Mod": np.fmod if attrs.get("fmod", 0) else np.mod, }[op_name] expected_value = op(lhs, rhs).astype(dtype_str) @@ -697,6 +697,31 @@ def test_mod(int_mode: bool): verify_binary_scalar("Mod", attrs={"fmod": fmod}, dtype=dtype) +@pytest.mark.parametrize( + "fmod, dtype, a_vals, b_vals", + [ + (0, TensorProto.INT32, [-5, 5, -5, 5], [3, 3, -3, -3]), + (1, TensorProto.INT32, [-5, 5, -5, 5], [3, 3, -3, -3]), + (1, TensorProto.FLOAT, [-5.5, 5.5, -5.5, 5.5], [3.0, 3.0, -3.0, -3.0]), + ], +) +def test_mod_constant_fold_negative_operands(fmod, dtype, a_vals, b_vals): + """Mod over two constants is folded at import time. The folded value must + match onnxruntime for negative operands, where integer mod (sign follows + divisor) and fmod (sign follows dividend) disagree.""" + a = make_constant_node("a", dtype, [4], a_vals) + b = make_constant_node("b", dtype, [4], b_vals) + mod_node = helper.make_node("Mod", ["a", "b"], ["c"], fmod=fmod) + graph = helper.make_graph( + [a, b, mod_node], + "mod_constant_fold_test", + inputs=[], + outputs=[helper.make_tensor_value_info("c", dtype, [4])], + ) + model = helper.make_model(graph, producer_name="mod_constant_fold_test") + check_correctness(model) + + SHAPE_PARAMS = [ ([[32, 32], [32, 32]], [32, 32]), ([[32, 1], [1, 2]], [32, 2]),