Function _zero_kernel

tinygrad/nn/__init__.py:324–326

(out:UOp)

  # TODO: how do we remove this dumb kernel and use Tensor.zeros?
  def _zero_kernel(out:UOp) -> UOp:
    i = UOp.range(out.numel(), 0)
    return out.flatten()[i].store(0).end(i).sink(arg=KernelInfo(name="zero"))
  grad_weight_uop = grad_weight_uop.custom_kernel(fxn=_zero_kernel)[0]

  # TODO: do we have a universal helper for this?