(inp, tuple_val, *, allow_newaxis=True)
| 134 | |
| 135 | |
| 136 | def unpack_getitem(inp, tuple_val, *, allow_newaxis=True): |
| 137 | assert isinstance(inp, Tensor) |
| 138 | if not isinstance(tuple_val, tuple): |
| 139 | tuple_val = (tuple_val,) |
| 140 | |
| 141 | def as_tensor(v): |
| 142 | if not isinstance(v, Tensor): |
| 143 | vi = np.ascontiguousarray(v, dtype=np.int32) |
| 144 | assert np.abs(vi - v).max() == 0, "bad index: {!r}".format(v) |
| 145 | v = Tensor(vi) |
| 146 | return v |
| 147 | |
| 148 | new_axes = [] |
| 149 | tensors = [] |
| 150 | items = [] |
| 151 | cur_axis = -1 |
| 152 | for i_idx, i in enumerate(tuple_val): |
| 153 | cur_axis += 1 |
| 154 | if i is np.newaxis: |
| 155 | if cur_axis >= 0: |
| 156 | new_axes.append(cur_axis) |
| 157 | continue |
| 158 | |
| 159 | if i is Ellipsis: |
| 160 | cur_axis = -1 |
| 161 | for j in tuple_val[:i_idx:-1]: |
| 162 | if j is Ellipsis: |
| 163 | raise IndexError("only one ellipsis is allowed") |
| 164 | if j is np.newaxis: |
| 165 | new_axes.append(cur_axis) |
| 166 | cur_axis -= 1 |
| 167 | continue |
| 168 | |
| 169 | item = [ |
| 170 | cur_axis, |
| 171 | ] |
| 172 | |
| 173 | def push(v, item, tensors): |
| 174 | if v is None: |
| 175 | item.append(False) |
| 176 | else: |
| 177 | item.append(True) |
| 178 | tensors.append(as_tensor(v)) |
| 179 | |
| 180 | if isinstance(i, slice): |
| 181 | if i.start is None and i.stop is None and i.step is None: |
| 182 | continue |
| 183 | push(i.start, item, tensors) |
| 184 | push(i.stop, item, tensors) |
| 185 | push(i.step, item, tensors) |
| 186 | item.append(False) # idx |
| 187 | else: |
| 188 | item += [False,] * 3 # begin, end, stop |
| 189 | push(i, item, tensors) |
| 190 | assert len(item) == 5 |
| 191 | items.append(item) |
| 192 | if new_axes: |
| 193 | raise IndexError("newaxis is not allowed here") |
no test coverage detected