diff --git a/test/test_ops.py b/test/test_ops.py index 98fc1912b745b..c1c49d2613229 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1236,6 +1236,7 @@ def test_slice_errors(self): with self.assertRaisesRegex(IndexError, "out of bounds"): a[1, -4] with self.assertRaisesRegex(IndexError, "single ellipsis"): a[..., ...] # IndexError: only single ellipsis with self.assertRaises(ValueError): a[::0, 1] # no 0 strides + with self.assertRaises(TypeError): a[:Tensor([3]), 1] # Tensor can't be used as a slice parameter with self.assertRaises(IndexError): b[:] # slice cannot be applied to a 0-dim tensor def test_slice_ellipsis(self): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 50da24a6859b7..27b4718c49420 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1063,6 +1063,9 @@ def _getitem(self, indices, v: Optional[Tensor] = None) -> Tensor: indices_filtered[dim] = ((index, index+1), 1) if index >= 0 else ((size+index, size+index+1), 1) for dim in type_dim[slice]: if (index := indices_filtered[dim]).step == 0: raise ValueError(f"{index=} on {dim=} cannot have 0 as step") + if not all(isinstance(x, (int, type(None))) for x in (index.start, index.stop, index.step)): + raise TypeError(f"Unsupported slice for dimension {dim}. Expected slice with integers or None, got slice(" + f"{', '.join(type(x).__name__ for x in (index.start, index.stop, index.step))}).") s, e, st = index.indices(self.shape[dim]) indices_filtered[dim] = ((0, 0) if (st * (e - s)) < 0 else (s, e) if st > 0 else (e+1, s+1), st) # skip all Tensor dims for basic indexing