nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

add kwargs to group (#17484)

Browse Source
This commit is contained in:
Raine
2026-08-10 17:04:38 -07:00
committed by GitHub
parent 7edf80a48a
commit 115bf9940f
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/uop/ops.py
+2 -2
View File
@@ -568,9 +568,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
in_tuple = self.src[0] if self.op is Ops.FUNCTION else self
assert in_tuple.op is Ops.TUPLE, f"gettuple requires FUNCTION or TUPLE source, got {self.op}"
return UOp(Ops.GETTUPLE, src=(self,), arg=idx)
def group(*srcs:UOp|None): # pylint: disable=no-self-argument
def group(*srcs:UOp|None, **kwargs): # pylint: disable=no-self-argument
if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]))
return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
def index(self, *srcs:UOp|int|None, **kwargs):
new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]