Function test_matmul

tests/python/relax/test_transform_gradient_numeric.py:198–231

(target, dev)

@tvm.testing.parametrize_targets("llvm")
def test_matmul(target, dev):
    @tvm.script.ir_module
    class Before:
        @R.function
        def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
            with R.dataflow():
                lv1 = R.matmul(x, y)
                lv2 = R.permute_dims(x)
                lv3 = R.matmul(lv2, y)
                lv4 = R.permute_dims(y)
                lv5 = R.matmul(x, lv4)
                lv6 = R.permute_dims(x)
                lv7 = R.permute_dims(y)
                lv8 = R.matmul(lv6, lv7)
                lv9 = lv1 + lv3 + lv5 + lv8
                gv = R.sum(lv9)
                R.output(gv)
            return gv

    After = relax.transform.Gradient("main")(Before)
    args = []
    for arg in After["main_adjoint"].params:
        shape = [int(l) for l in arg.struct_info.shape]
        args.append(rand("float32", *shape))

    vm_before = _legalize_and_build(Before, target, dev)
    vm_after = _legalize_and_build(After, target, dev)
    _, grad = vm_after["main_adjoint"](*args)

    def func(*inputs):
        loss = vm_before["main"](*[tvm.runtime.tensor(i) for i in inputs])
        return loss.numpy()

    check_numerical_grads(func, [i.numpy() for i in args], [i.numpy() for i in grad])


if __name__ == "__main__":