Function test_emit_te

tests/python/relax/test_transform_gradient_te_register.py:121–145

(register_te_grads)

def test_emit_te(register_te_grads):
    # Build the target module using emit_te
    def f_mul(src1, src2):
        def mul(*idx):
            return src1[idx] * src2[idx]

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

    a = relax.Var("a", relax.TensorStructInfo([5, 5], "float32"))
    b = relax.Var("b", relax.TensorStructInfo([5, 5], "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_1())


def test_call_tir(register_te_grads):