()
| 66 | |
| 67 | |
| 68 | def test_implicit_op(): |
| 69 | m, n = tvm.tirx.Var("m", "int64"), tvm.tirx.Var("n", "int64") |
| 70 | x = rx.Var("x", R.Tensor([m, n], "float32")) |
| 71 | y = rx.Var("y", R.Tensor([m, n], "float32")) |
| 72 | func = rx.Var( |
| 73 | "func", |
| 74 | R.Callable( |
| 75 | [R.Tensor([m, n], "float32")], |
| 76 | R.Callable( |
| 77 | [R.Tensor([m, n], "float32")], |
| 78 | R.Tuple, |
| 79 | ), |
| 80 | ), |
| 81 | ) |
| 82 | |
| 83 | def _check_call(expr, op_name: str): |
| 84 | assert isinstance(expr, rx.Call) |
| 85 | if not op_name.startswith("relax."): |
| 86 | op_name = "relax." + op_name |
| 87 | op = tvm.ir.Op.get(op_name) |
| 88 | assert expr.op == op |
| 89 | |
| 90 | # Comparison operators |
| 91 | _check_call(x > y, "greater") |
| 92 | _check_call(x >= y, "greater_equal") |
| 93 | _check_call(x < y, "less") |
| 94 | _check_call(x <= y, "less_equal") |
| 95 | |
| 96 | # Arithmetic operators |
| 97 | _check_call(-x, "negative") |
| 98 | _check_call(x + y, "add") |
| 99 | _check_call(x - y, "subtract") |
| 100 | _check_call(x * y, "multiply") |
| 101 | _check_call(x / y, "divide") |
| 102 | _check_call(x // y, "floor_divide") |
| 103 | _check_call(x**y, "power") |
| 104 | # _check_call(x % y, "mod") <= relax.mod is not implemented yet |
| 105 | |
| 106 | # Cast |
| 107 | _check_call(x.astype("float32"), "astype") |
| 108 | |
| 109 | # Call |
| 110 | call_expr = func(y)(y) |
| 111 | assert isinstance(call_expr.op, rx.Call) |
| 112 | assert call_expr.op.op == func |
| 113 | |
| 114 | # GetTupleItem |
| 115 | ## Eager get item for tuple |
| 116 | tuple_expr = rx.Tuple((x, y)) |
| 117 | assert tuple_expr[0] == x |
| 118 | assert tuple_expr[1] == y |
| 119 | |
| 120 | ## Eager get item for ShapeExpr |
| 121 | shape_expr = rx.ShapeExpr((1, 2)) |
| 122 | assert shape_expr[0] == 1 |
| 123 | assert shape_expr[1] == 2 |
| 124 | |
| 125 | ## Create TupleGetItem for other expr |
nothing calls this directly
no test coverage detected
searching dependent graphs…