From 2193d0edfa41d9faa95f8ad4b682c7b44e0eba98 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Thu, 5 Feb 2026 12:03:45 +0800
Subject: [PATCH] fix arg order

---
 tinygrad/gradient.py |  2 +-
 tinygrad/tensor.py   |  2 +-
 tinygrad/uop/ops.py  | 11 ++++++++++-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/tinygrad/gradient.py b/tinygrad/gradient.py
index e6e48908f5..b5b9a4f904 100644
--- a/tinygrad/gradient.py
+++ b/tinygrad/gradient.py
@@ -14,7 +14,7 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
 def call_gradient(ctx:UOp, k:UOp):
-  if k.arg is not None: return (None,) + k.arg(ctx, k)
+  if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
   # auto-differentiate the function
   fxn, args = k.src[0], k.src[1:]
   params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index cb9e4c0dbc..eec5eeac41 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -240,7 +240,7 @@ class Tensor(OpMixin):
       param = UOp.param(slot, self.dtype, self.shape, self.device)
     return Tensor(param, device=self.device)
   def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
-    return Tensor(UOp.call(*[t.uop for t in (self,)+lst], fxn=fxn.uop if isinstance(fxn, Tensor) else fxn, arg=grad_fxn), device=self.device)
+    return Tensor((fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn), device=self.device)
 
   def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
     """
diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 1a1ddbeced..06b24d1a5d 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -818,7 +818,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     src = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + (() if device is None else (UOp(Ops.DEVICE, arg=device),))
     return UOp(Ops.PARAM, dtype, src, arg=slot)
 
-  def call(*srcs:UOp, fxn:UOp, arg:Any|None) -> UOp: return UOp(Ops.CALL, fxn.dtype, (fxn,)+srcs, arg)
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> UOp:
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
     contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
     kernel = UOp(Ops.CUSTOM_KERNEL, src=contig_srcs, arg=CustomKernel(fxn=fxn, grad_fxn=grad_fxn))
@@ -843,6 +844,14 @@ class CustomKernel:
   def __reduce__(self): return (CustomKernel, (panic,))
   def __repr__(self): return f"CustomKernel({id(self.fxn)})"
 
+@dataclass(frozen=True)
+class CallInfo:
+  grad_fxn: Callable|None = None
+  metadata: tuple[Metadata, ...] = ()
+  # CallInfo can't be pickled or reconstructed as a str when grad_fxn is set
+  def __reduce__(self): return (CallInfo, ())
+  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata})"
+
 @dataclass(frozen=True)
 class Kernel:
   ast: UOp