Skip to content

Fix product reduction gradient when an element is zero - #528

Open
lnuic wants to merge 1 commit into
masterfrom
fix-prod-grad-zero
Open

Fix product reduction gradient when an element is zero#528
lnuic wants to merge 1 commit into
masterfrom
fix-prod-grad-zero

Conversation

@lnuic

@lnuic lnuic commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

The gradient of a product should be the product of the other entries, but we
computed it as result / value, which breaks when an entry is zero.
dr.prod caught the divide by zero and returned 0, which is the wrong answer: for
[2, 0, 4] the gradient should be [0, 8, 0] and previous result was [0, 0, 0].

The axis/block version didn't check at all and returned NaN.

The cost is two extra reductions in the backward pass, so a product backward is about 3x slower:

backend op N before after
LLVM prod 16M 3.8 ms 13.0 ms
LLVM block_reduce 16M 6.0 ms 13.0 ms
CUDA prod 16M 0.044 ms 0.135 ms
CUDA block_reduce 16M 0.090 ms 0.314 ms

That could be brought back to a single pass by combining the two reductions, but
it needs a new reduction primitive in drjit-core, which felt like too much for
this one fix. I am happy to add that if needed.

If the slowdown isn't worth it I am happy to leave the gradient approximate and
just guard the block version so it returns 0 instead of NaN.

@wjakob

wjakob commented Jul 28, 2026

Copy link
Copy Markdown
Member

Hi Lovro -- are you sure about needing new backend operations? Dr.Jit supports block_prefix_reduce which seems to be just what is needed. What I don't like about this version is that it seems to make the primal slower even when no derivatives are ultimately needed.

@wjakob
wjakob force-pushed the master branch 7 times, most recently from 60b72ec to 23ef9cd Compare August 22, 2026 03:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants