()
| 102 | |
| 103 | |
| 104 | def test_runtime_indexing(): |
| 105 | def data_gen(): |
| 106 | return [ |
| 107 | np.float32([[0, 1, 2], [3, 4, 5]]), |
| 108 | np.float32([[0, 1], [2, 3], [4, 5]]), |
| 109 | ] |
| 110 | |
| 111 | src = fn.external_source(data_gen) |
| 112 | lo_idxs = [np.array(x, dtype=np.int64) for x in [1, -5, 0, 2, -2, 1]] |
| 113 | hi_idxs = [np.array(x, dtype=np.int16) for x in [5, -1, 1, 2, 4]] |
| 114 | lo0 = fn.external_source(source=lo_idxs, batch=False, cycle=True) |
| 115 | hi1 = fn.external_source(source=hi_idxs, batch=False, cycle=True) |
| 116 | pipe = index_pipe(src, lambda x: x[lo0:, :hi1]) |
| 117 | j = 0 |
| 118 | k = 0 |
| 119 | for _ in range(4): |
| 120 | inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run()) |
| 121 | for i in range(len(inp)): |
| 122 | x = inp.at(i) |
| 123 | # fmt: off |
| 124 | ref = x[lo_idxs[j]:, :hi_idxs[k]] |
| 125 | # fmt: on |
| 126 | j = (j + 1) % len(lo_idxs) |
| 127 | k = (k + 1) % len(hi_idxs) |
| 128 | assert np.array_equal(ref, cpu.at(i)) |
| 129 | assert np.array_equal(ref, gpu.at(i)) |
| 130 | |
| 131 | |
| 132 | def test_runtime_stride_dim1(): |
nothing calls this directly
no test coverage detected