Function set_input_trial

tests/python/relax/test_vm_build.py:1063–1093

(vm: relax.VirtualMachine, device: tvm.runtime.Device)

def set_input_trial(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None:
    a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    vm.set_input("main", a, b)
    vm.invoke_stateful("main")
    res0 = vm.get_outputs("main")

    data_dict = {"x": a, "w": b}
    vm.set_input("main", **data_dict)
    vm.invoke_stateful("main")
    res1 = vm.get_outputs("main")
    tvm.testing.assert_allclose(res0.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)
    tvm.testing.assert_allclose(res0.numpy(), res1.numpy(), rtol=1e-7, atol=1e-7)

    # bug! If you don't bind the Tensor to a var, the memory will get corrupted.
    # Possibly due to object lifecycles and other FFI issues
    a = tvm.runtime.tensor(np.array(2).astype("int32"), device)
    vm.set_input("test_vm_tuple", a)
    vm.invoke_stateful("test_vm_tuple")
    res2 = vm.get_outputs("test_vm_tuple")
    # the results are Tensors wrapped around scalars,
    # so we have to get the scalar out of the Tensor
    assert tuple(map(lambda a: int(a.numpy()), res2)) == (2, 2)

    b = tvm.runtime.tensor(np.array(1).astype("int32"), device)
    vm.set_input("test_vm_nested_tuple", b)
    vm.invoke_stateful("test_vm_nested_tuple")
    res3 = vm.get_outputs("test_vm_nested_tuple")
    assert len(res3) == 2 and len(res3[0]) == 2 and len(res3[0][1]) == 1
    result_cast = ((int(res3[0][0].numpy()), (int(res3[0][1][0].numpy()),)), int(res3[1].numpy()))
    assert result_cast == ((1, (1,)), 1)


def set_input_attempt_stateless(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None: