Integer-slicing of missing dimensions
Example:
>>> math.meshgrid(x=4, y=3).a[0]
(x=4, y=3, vector=2) int32  0 < ... < 3
Proper slicing (start,stop,step) is not supported as the type of the dimension is unknown.
@@ -414,9 +414,12 @@ class NativeTensor(Tensor):
new_shape = self.shape
selections = [slice(None)] * self.rank
for name, selection in selection.items():
+ if name in self.shape:
selections[self.shape.index(name)] = selection
if isinstance(selection, int):
new_shape = new_shape.without(name)
+ else:
+ assert isinstance(selection, int), f"Attempting slice missing dimension {name} with {selection}"
gathered = self.tensor[tuple(selections)]
new_shape = new_shape.with_sizes(native_math.staticshape(gathered))
return NativeTensor(gathered, new_shape)
