From 357dac8425b57fa8db108f4b609b4e7692168cc7 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Sun, 19 Oct 2025 19:11:05 -0700 Subject: [PATCH] feat: allow tuple indexing on uops (#12797) --- tinygrad/uop/ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 55ee4a69c1..488fc0774c 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -324,7 +324,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): def detach(self): return UOp(Ops.DETACH, self.dtype, (self,)) def index(self, *srcs:UOp|None, **kwargs): return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs) - def __getitem__(self, idx): return self.index(idx) + def __getitem__(self, *idx): return self.index(*idx) def const_like(self, b:ConstLike): # constants can optionally have a DEVICE source return UOp.const(self.dtype, b, device=self._device, shape=self._shape)