| 57 | |
| 58 | @tvm.register_global_func("rpc.test.remote_return_nd") |
| 59 | def _my_module(name): |
| 60 | # Use closure to check the ref counter correctness |
| 61 | nd = tvm.runtime.tensor(np.zeros(10).astype("float32")) |
| 62 | |
| 63 | if name == "get_arr": |
| 64 | return lambda: nd |
| 65 | if name == "ref_count": |
| 66 | # Imported lazily: rpc.server imports this module to register the |
| 67 | # rpc.test.* helpers, so a top-level ``import tvm.testing`` would drag |
| 68 | # tvm.testing (which imports pytest) into the plain ``import tvm`` path. |
| 69 | from tvm.testing import object_use_count |
| 70 | |
| 71 | return lambda: object_use_count(nd) |
| 72 | if name == "get_elem": |
| 73 | return lambda idx: nd.numpy()[idx] |
| 74 | if name == "get_arr_elem": |
| 75 | return lambda arr, idx: arr.numpy()[idx] |
| 76 | raise RuntimeError("unknown name") |