tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 11:56:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

clean up reduce MUL gradient (#17447)

Browse Source
This commit is contained in:
chenyu
2026-08-07 10:08:04 -04:00
committed by GitHub
parent 0c96cdc300
commit fca695a36f
1 changed files with 3 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/mixin/gradient.py
+3 -5
View File
@@ -8,12 +8,10 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,)
if op == Ops.MUL:
x, axes = ret.src[0], tuple(range(ret.arg[1]))
is_zero = x.eq(0)
safe_x = is_zero.where(1, x)
# d(prod x)/dx_j = prod_{i!=j} x_i: ret/x_j whenever x_j != 0 (any zero makes ret 0), else the product of the others
safe_x, axes = (is_zero:=(x:=ret.src[0]).eq(0)).where(1, x), tuple(range(ret.arg[1]))
zero_count = is_zero.cast(sum_acc_dtype(is_zero.dtype))._rop(Ops.ADD, axes)
zero_grad = (is_zero & zero_count.eq(1)).where(safe_x._rop(Ops.MUL, axes), 0)
return (ctx * zero_count.eq(0).where(ret / safe_x, zero_grad),)
return (ctx * is_zero.where(zero_count.eq(1).where(safe_x._rop(Ops.MUL, axes), 0), ret/safe_x),)
def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
"""Remove unused PARAMs from body and return compacted (body, args)."""