Function test_op_correctness

tests/python/relax/test_op_binary.py:29–61

()

def test_op_correctness():
    x = relax.Var("x", R.Tensor((2, 3), "float32"))
    y = relax.Var("y", R.Tensor((2, 3), "float32"))
    assert relax.op.add(x, y).op == Op.get("relax.add")
    assert relax.op.divide(x, y).op == Op.get("relax.divide")
    assert relax.op.floor_divide(x, y).op == Op.get("relax.floor_divide")
    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.atan2(x, y).op == Op.get("relax.atan2")
    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")
    assert relax.op.greater_equal(x, y).op == Op.get("relax.greater_equal")
    assert relax.op.less(x, y).op == Op.get("relax.less")
    assert relax.op.less_equal(x, y).op == Op.get("relax.less_equal")
    assert relax.op.not_equal(x, y).op == Op.get("relax.not_equal")

    x = relax.Var("x", R.Tensor((2, 3), "int32"))
    y = relax.Var("y", R.Tensor((2, 3), "int32"))
    assert relax.op.bitwise_and(x, y).op == Op.get("relax.bitwise_and")
    assert relax.op.bitwise_or(x, y).op == Op.get("relax.bitwise_or")
    assert relax.op.bitwise_xor(x, y).op == Op.get("relax.bitwise_xor")
    assert relax.op.left_shift(x, y).op == Op.get("relax.left_shift")
    assert relax.op.right_shift(x, y).op == Op.get("relax.right_shift")

    x = relax.Var("x", R.Tensor((2, 3), "bool"))
    y = relax.Var("y", R.Tensor((2, 3), "bool"))
    assert relax.op.logical_and(x, y).op == Op.get("relax.logical_and")
    assert relax.op.logical_or(x, y).op == Op.get("relax.logical_or")
    assert relax.op.logical_xor(x, y).op == Op.get("relax.logical_xor")


def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: relax.StructInfo):