- for i in range(index.dim()):
- v = s_index[i]
- j = s_dst.index(v)
- a = (
- torch.arange(index.size(i))
- .reshape((1,) * j + (-1,) + (1,) * (len(s_dst) - j - 1))
- .expand(shape)
- )
- idx.append(a)
+ for i in range(x.dim()):
+ v = s_x[i]
+ if v == "_":
+ idx.append(do(indexes[n_index], s_indexes[n_index]))
+ n_index += 1
+ else:
+ j = s_dst.index(v)
+ a = (
+ torch.arange(x.size(i))
+ .reshape((1,) * j + (-1,) + (1,) * (len(s_dst) - j - 1))
+ .expand(shape)
+ )
+ idx.append(a)