Function check_saved_func

python/tvm/relax/testing/vm.py:78–87

(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any])

def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any]) -> Object:
    # uses save_function to create a closure with the given inputs
    # and ensure the result is the same
    # (assumes the functions return tensors and that they're idempotent)
    saved_name = f"{func_name}_saved"
    vm.save_function(func_name, saved_name, *inputs)
    res1 = vm[func_name](*inputs)
    res2 = vm[saved_name]()
    tvm.testing.assert_allclose(res1.numpy(), res2.numpy(), rtol=1e-7, atol=1e-7)
    return res1


@tvm.register_global_func("test.vm.check_if_defined")