diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py index 8dd4eff5c703..e9bc9a7a3e98 100644 --- a/python/tvm/relax/expr.py +++ b/python/tvm/relax/expr.py @@ -185,8 +185,7 @@ def __rfloordiv__(self, other: Expr) -> "ExprWithOp": return _binary_rhs_helper(other) def __mod__(self, other: Expr) -> "ExprWithOp": - # TODO(siyuan): Support it after mod operator is supported in relax - raise ValueError("relax.mod is not supported yet.") + return _binary_op_helper(self, other, _op_ffi_api.mod) # type: ignore def __rmod__(self, other: Expr) -> "ExprWithOp": return _binary_rhs_helper(other) diff --git a/tests/python/relax/test_op_binary.py b/tests/python/relax/test_op_binary.py index 20c111495d6a..3376569bf349 100644 --- a/tests/python/relax/test_op_binary.py +++ b/tests/python/relax/test_op_binary.py @@ -33,6 +33,8 @@ def test_op_correctness(): assert relax.op.multiply(x, y).op == Op.get("relax.multiply") assert relax.op.power(x, y).op == Op.get("relax.power") assert relax.op.subtract(x, y).op == Op.get("relax.subtract") + assert relax.op.mod(x, y).op == Op.get("relax.mod") + assert relax.op.floor_mod(x, y).op == Op.get("relax.floor_mod") assert relax.op.equal(x, y).op == Op.get("relax.equal") assert relax.op.greater(x, y).op == Op.get("relax.greater") @@ -70,6 +72,8 @@ def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: r (relax.op.subtract, tir.Sub), (relax.op.maximum, tir.Max), (relax.op.minimum, tir.Min), + (relax.op.mod, tir.Mod), + (relax.op.floor_mod, tir.FloorMod), )