geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

feat: allow tuple indexing on uops (#12797)

Browse Source
This commit is contained in:
wozeparrot
2025-10-19 19:11:05 -07:00
committed by GitHub
parent ba593f7b98
commit 357dac8425
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/uop/ops.py
+1 -1
View File
@@ -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)