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]