From 115bf9940fa3236a2949c101e95959cbb3f37e2e Mon Sep 17 00:00:00 2001 From: Raine Date: Mon, 10 Aug 2026 21:04:38 -0300 Subject: [PATCH] add kwargs to group (#17484) --- tinygrad/uop/ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 9ae599b987..723e15fdff 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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]