google/jax
copybara-service[bot]'s Description of Work
[JAX] Fix jax.scipy.linalg.toeplitz (and hankel) NaN propagation
Problem:
In https://github.com/jax-ml/jax/issues/38636 (and https://github.com/jax-ml/jax/issues/38635), jax.scipy.linalg.toeplitz and jax.scipy.linalg.hankel propagated NaNs incorrectly because their underlying implementations used lax.conv_general_dilated_patches. This convolution implementation expands NaN pollution across the entire output matrix due to structural zero padding interacting with NaN.
Solution:
Replace lax.conv_general_dilated_patches with lax.broadcasted_iota indexing in _toeplitz and _hankel. This propagates NaNs only to their respective indexed locations. Added unit tests with nonfinite values to ensure correctness.
Fixes https://github.com/jax-ml/jax/issues/38635
Fixes https://github.com/jax-ml/jax/issues/38636