google/jax
Pull Request Overview
- Opened on October 21, 2026
- Status Merged
- Commit count 1 with first commit October 22, 2026
Total Delta
Open Days
Test Delta
How long has this pull request spent in each phase of its lifecycle?
Data pending calculation for pull request
[JAX] Fix jax.scipy.linalg.toeplitz (and hankel) NaN propagation
[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
PR was closed without comments.