Function test_tir_var

tests/python/relax/test_transform_gradient_te_register.py:358–383

(register_te_grads)

def test_tir_var(register_te_grads):
    def f_mul(src1, src2):
        def mul(*idx):
            return src1[idx] * src2[idx]

        return tvm.te.compute(src1.shape, mul, name="f_mul")

    n = tirx.Var("n", "int64")
    a = relax.Var("a", relax.TensorStructInfo([n, n], "float32"))
    b = relax.Var("b", relax.TensorStructInfo([n, n], "float32"))

    bb = relax.BlockBuilder()
    with bb.function("main", [a, b]):
        with bb.dataflow():
            d = bb.emit(
                bb.call_te_with_grad(
                    f_mul, a, b, primfunc_name_hint="f_mul", te_grad_name="f_mul_grad"
                )
            )
            out = bb.emit_output(R.sum(d))
        bb.emit_func_output(out)

    Before = bb.get()
    After = Gradient("main")(Before)
    assert_structural_equal(After, get_expected_3())
    relax.analysis.well_formed(After)


if __name__ == "__main__":