Function test_emit_te_kwargs

tests/python/relax/test_transform_gradient_te_register.py:229–254

(register_te_grads)

def test_emit_te_kwargs(register_te_grads):
    # Build the target module using emit_te
    def f_mul2(src):
        return tvm.te.compute(src.shape, lambda *idx: src[idx] * T.float32(2), name="f_mul2")

    a = relax.Var("a", relax.TensorStructInfo([5, 5], "float32"))

    bb = relax.BlockBuilder()
    with bb.function("main", [a]):
        with bb.dataflow():
            d = bb.emit(
                bb.call_te_with_grad(
                    f_mul2,
                    a,
                    primfunc_name_hint="f_mul",
                    te_grad_name="f_mulk_grad",
                    te_grad_kwargs={"k": T.float32(2)},
                )
            )
            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_2())


def test_call_tir_kwargs(register_te_grads):