()
| 153 | |
| 154 | |
| 155 | def test_vm_shapeof(): |
| 156 | ib = relax.ExecBuilder() |
| 157 | shape = (32, 16) |
| 158 | arr = tvm.runtime.tensor(np.random.rand(*shape)) |
| 159 | with ib.function("main", num_inputs=0): |
| 160 | ib.emit_call("vm.builtin.shape_of", args=[arr], dst=ib.r(0)) |
| 161 | ib.emit_ret(ib.r(0)) |
| 162 | ex = ib.get() |
| 163 | vm = relax.VirtualMachine(ex, tvm.cpu()) |
| 164 | res = vm["main"]() |
| 165 | for i, s in enumerate(res): |
| 166 | assert s == shape[i] |
| 167 | |
| 168 | |
| 169 | def test_vm_storage(): |