fix(autodiff): zero-safe cumprod gradient (no division by input)

[AnayGarodia]

Fixes #3864.

Problem

The cumprod backward pass computed reverse_cumsum(grad * output) / input, dividing by the input and producing NaN gradients whenever the input contains zeros. Four autodiff tests for this case were #[ignore]d with a TODO pointing at this issue.

Fix

Replace the division with the standard zero-safe formula grad_input[i] = left[i] * tail[i], where:

• left[i] = prod(input[0..i]) — exclusive prefix product (built by prepending 1 along dim, dropping the last element, then an inclusive cumprod)
• tail[i] = grad[i] + input[i+1] * tail[i+1] — reverse accumulation along dim

This never divides by input[i], so it stays finite at zeros and matches torch.cumprod's backward.

Testing

Un-ignored the four zero-input tests (zero_in_middle, zero_at_start, zero_at_end, multiple_zeros), whose expected gradients are taken from PyTorch — they now pass.

cargo test -p burn-backend-tests --features ndarray,std --test autodiff cumprod   # 12 passed, 0 ignored
cargo clippy -p burn-autodiff --all-targets -- -D warnings                        # clean
cargo fmt --check                                                                 # clean

Verified numerically against PyTorch references, e.g. input [2,0,3,4] → grad [1,32,0,0] (was NaN). The fix uses only backend-agnostic tensor ops; I ran it against the ndarray backend.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Charles23R]

How does the performance compare with the previous implementation on GPU? This is probably way slower. We should at the very least keep the original, faster implementation if there are no 0s in the input, and keep a note to reimplement with a scan/fold operation eventually like mentioned in #5141

[AnayGarodia]

Good call on the perf question — the sequential recurrence was indeed the weak point, and benchmarking it on wgpu/Metal surfaced a second problem: a few hundred chained slice/mul/add ops made the fusion pass fall over (repeated elemwise_fuse shader-generation failures, then a stack overflow at n=1024).

So I dropped the recurrence entirely and reworked the gradient into the masked parallel form PyTorch's cumprod_backward uses. Each lane is split at its first zero along dim:

• before the first zero: the classic reverse_cumsum(grad * output) / input is exact (every divisor is nonzero),
• at the first zero: reverse_cumsum(grad * cumprod(input with that zero set to 1)) evaluated at that position,
• after it: exactly zero.

The regions are computed with cumsum/cumprod plus elementwise masks, so it's a constant number of kernel launches for every input — no data-dependent branching, no readback, no O(n) loop, and rows with no zeros reduce to the classic formula. Multiple zeros per lane are handled (everything past the first zero is zero).

Perf on an Apple M-series GPU (Metal via wgpu, fwd+bwd+sync, 50 iters — noisy laptop numbers, treat as ballpark):

shape	old (main)	this branch
32x128, no zeros	5.5 ms	2.9 ms
32x1024, no zeros	3.5 ms	9.7 ms
8x4096, no zeros	6.1 ms	9.2 ms
32x1024, with zeros	6.6 ms (NaN grads)	3.4 ms (correct)

Within noise on this machine, both are flat in n; the new form adds ~8 small elementwise kernels of constant overhead. A scan/fold primitive would still be nice for other ops, but this no longer depends on it. All 12 cumprod autodiff tests pass (zero-at-start/middle/end, multiple zeros, checkpointing).