[math] Preserve item names when slicing with int
@@ -322,9 +322,12 @@ class Tensor:
else:
from ._magic_ops import stack
result = [sliced[{dim: i}] for i in selection_int]
- item_names = [str(n) for n in selection]
- stack_dim = self.shape[dim] if dim in self.shape else channel(dim)
- sliced = stack({n: r for n, r in zip(item_names, result)}, stack_dim.with_size(None))
+ if all(isinstance(n, str) for n in selection) or self.shape.get_item_names(dim) is not None:
+ item_names = [i if isinstance(i, str) else self.shape.get_item_names(dim)[i] for i in selection]
+ stack_dim = self.shape[dim].with_size(len(selection))._with_item_names((tuple(item_names),)) if dim in self.shape else channel(**{dim: item_names})
+ else:
+ stack_dim = self.shape[dim].with_size(len(selection)) if dim in self.shape else channel(dim)
+ sliced = stack(result, stack_dim)
elif isinstance(selection, Tensor) and selection.dtype.kind == bool:
from ._ops import boolean_mask
sliced = boolean_mask(sliced, dim, selection)
