Function _compact_params

tinygrad/gradient.py:17–20

Remove unused PARAMs from body and return compacted (body, args).

(body:UOp, all_args:tuple[UOp, ...])

  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)

def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
  """Remove unused PARAMs from body and return compacted (body, args)."""
  used = sorted({p.arg: p for p in body.toposort() if p.op is Ops.PARAM}.items())
  return body.substitute({p: p.replace(arg=j) for j,(_, p) in enumerate(used)}, walk=True), tuple(all_args[i] for i,_ in used)

def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
  fxn, args = k.src[0], k.src[1:]