Indexing utils#
Index a tensor along a given dimension using an index tensor, replacing the shape along the given dimension with the shape of the index tensor. |
- drtk.utils.indexing.index(x, idxs, dim)[source]#
Index a tensor along a given dimension using an index tensor, replacing the shape along the given dimension with the shape of the index tensor.
Example: x: [8, 7306, 3] idxs: [11000, 3]
y = index(x, idxs, dim=1) -> y: [8, 11000, 3, 3] with each y[b, i, j, k] = x[b, idxs[i, j], k]