DivToMul

import sys

sys.path.extend(["../../../tests"])
# from tools.tag_span import _create_span, _set_span, _verify_structural_equal_with_span
from tools.torch_utils import verify_model

import tvm
from tvm import relay
import numpy as np

def test_div_to_mul(dtype, rtol):
    x = relay.var("x", relay.TensorType((), dtype))
    y = relay.Constant(tvm.nd.array(np.array([1.5]).astype(dtype)))
    z = x / y
    mod = tvm.IRModule.from_expr(z)
    mod.show()
    transformed = relay.transform.DivToMul()(mod)
    transformed.show()
    assert transformed["main"].body.op.name == "multiply"
    np.testing.assert_allclose(transformed["main"].body.args[1].data.numpy()[0], 1 / 1.5, rtol=rtol)

for dtype, rtol in [("float16", 1e-3), ("float32", 1e-7), ("float64", 1e-12)]:
    test_div_to_mul(dtype, rtol)

def @main(%x: float16) {
  divide(%x, meta[relay.Constant][0])
}

def @main(%x: float16 /* ty=float16 */) -> Tensor[(1), float16] {
  multiply(%x, meta[relay.Constant][0] /* ty=Tensor[(1), float16] */) /* ty=Tensor[(1), float16] */
}

def @main(%x: float32) {
  divide(%x, meta[relay.Constant][0])
}

def @main(%x: float32 /* ty=float32 */) -> Tensor[(1), float32] {
  multiply(%x, meta[relay.Constant][0] /* ty=Tensor[(1), float32] */) /* ty=Tensor[(1), float32] */
}

def @main(%x: float64) {
  divide(%x, meta[relay.Constant][0])
}

def @main(%x: float64 /* ty=float64 */) -> Tensor[(1), float64] {
  multiply(%x, meta[relay.Constant][0] /* ty=Tensor[(1), float64] */) /* ty=Tensor[(1), float64] */
}