Enhancement: Track zero propagation through reduce_prod

[bdrhill]

Summary

reduce_prod produces conservative sparsity when any input is known to be zero. If the product contains a zero, the entire product is zero with zero gradient for all inputs. Currently, all inputs are detected as contributing.

MWE

import jax.numpy as jnp
import asdex

def f(x):
    # Multiply by array with known zero
    return jnp.prod(x * jnp.array([1.0, 0.0, 1.0]))

x = jnp.array([2.0, 3.0, 4.0])
pattern = asdex.jacobian_sparsity(f, x)
print(pattern)
# SparsityPattern(1×3, nnz=2, sparsity=33.3%)
# ● ⋅ ●

# But x * [1,0,1] = [2, 0, 4], and prod([2,0,4]) = 0
# True Jacobian: [0. 0. 0.] (all zeros because product is 0)

Proposed Fix

In prop_reduce for reduce_prod:

1. Check if any input element is known to be exactly zero (via state_consts)
2. If a zero is present in the product, all gradients are zero → return empty sparsity

This extends the existing zero-tracking in prop_mul to reductions.

Edge Cases

• Product of positive numbers: no zeros, current behavior is correct
• Product containing zero: all gradients are zero
• Product containing zero multiplied by infinity: technically nan, but conservatively treat as zero-sparsity

Impact

Affects any computation involving products where zeros are introduced (e.g., masked operations, sparse structures). Relatively niche but follows the existing pattern of zero-tracking in mul.

[adrhill]

This only globally holds if there is a constant input that is zero.