JAX PR #40803: Why sinc Needed Exact Range Reduction Before Autodiff
JAX PR #40803 replaces the real-valued implementation of jnp.sinc with an additive range-reduction algorithm.
The mathematical definition is familiar:
sinc(x) = sin(pi × x) / (pi × x)Computing that expression literally is not numerically equivalent to evaluating the mathematical function.
Multiplying by a floating-point approximation of pi loses the exact location of every nonzero integer root. For sufficiently large finite inputs, pi × x can overflow even though sinc(x) should remain finite and close to zero. Repeated automatic differentiation also magnifies small polynomial and cancellation errors that may be harmless in the primal value.
The new implementation does not send the full input into sin(pi × x).
It writes each finite real input as:
x = n + rwhere:
n = round_to_even(x)
r = x - n
r is in [-0.5, 0.5]It then evaluates a higher-degree polynomial only on the small remainder and reconstructs the sign and scale of the original result.
The old branch still multiplied the full input by pi
Before this PR, the real and complex implementations shared a common structure.
They first formed:
u = pi × xFor small |u|, JAX evaluated a Taylor series for:
sin(u) / uFor larger values, it used the direct expression.
That local series avoided cancellation around zero. It did not solve the full-range argument-reduction problem.
Consider an exact integer input:
x = 1Mathematically:
sin(pi) = 0
sinc(1) = 0A floating-point constant named pi is not the exact transcendental value. The multiplication produces an approximation, and the hardware sine receives a number near pi rather than an exact symbolic multiple of pi.
The returned value can therefore be a small residual instead of the exact root.
The same problem appears near every nonzero integer. The sine implementation must recover a tiny remainder from an argument that already contains the rounding error introduced by the multiplication.
Large finite inputs could become NaN
The old implementation protected an input that was itself positive or negative infinity.
It did not protect every finite input for which:
pi × xoverflowed.
For a large finite float32, the multiplication can become infinity. The following direct path then attempts:
sin(infinity) / infinityand returns NaN.
That result does not reflect the mathematical behavior of sinc.
For large finite values:
|sinc(x)| <= 1 / (pi × |x|)so the function approaches zero.
The new implementation never needs to form pi × x for the full input. A huge representable floating-point number is already integral at its available precision, so rounding gives:
n = x
r = 0and the result can be returned as zero without overflowing an intermediate argument.
Additive range reduction keeps the important remainder
For finite real inputs, the patch computes:
n = round_to_even(x)
r = x - nBecause n is the nearest representable integer-valued floating-point number, the subtraction satisfies the conditions behind Sterbenz’s lemma. The remainder is obtained exactly for the finite inputs covered by this construction.
The trigonometric identity is:
sin(pi × (n + r))
=
(-1)^n × sin(pi × r)Therefore:
sinc(n + r)
=
(-1)^n × sinc(r) × r / xNow the trigonometric approximation sees only:
r in [-0.5, 0.5]instead of the full magnitude of x.
The algorithm preserves the exact integer center and evaluates only the local displacement from that center.
The real path no longer depends on a general sine call
For the reduced remainder, JAX expands:
sin(pi × r) / (pi × r)as an even polynomial in:
t = r²The expansion is:
1
- (pi² / 3!) × t
+ (pi⁴ / 5!) × t²
- ...The coefficients are evaluated with Horner’s rule.
For float32, the merged code uses six stored coefficients for the polynomial after the leading one.
For float64, it uses eleven.
These degrees are intentionally higher than the primal value alone would require. Automatic differentiation repeatedly differentiates the polynomial, removes powers of the variable, and multiplies coefficients by factorial terms. An approximation that is adequate for the function value can be less accurate for its fourth derivative.
The implementation therefore treats derivative accuracy as part of the numerical contract, not as an accidental property of the primal approximation.
stop_gradient keeps the interval choice discrete
The integer center is calculated as:
n = stop_gradient(round_to_even(x))Within one interval around an integer, n is a discrete range-reduction choice. It should not contribute a derivative.
The remainder is:
r = x - nso its derivative with respect to x remains one away from the rounding boundaries.
Without stopping the derivative through n, the autodiff graph would include the mechanics of choosing an interval center rather than only the smooth sinc function reconstructed within that interval.
The patch is not adding a separate hand-written derivative formula. It constructs a primal expression whose ordinary automatic derivatives remain stable.
Integer roots need a value and a derivative
At a nonzero integer:
r = 0The primal sinc value must be positive zero.
The first derivative is generally not zero:
sinc'(n) = (-1)^n / nA naive special case that replaces the entire branch with a constant zero would erase that derivative.
The patch uses this expression at the r == 0 boundary:
y - stop_gradient(y)Its primal value is zero because the two values cancel.
Its derivative is the derivative of y, because the stopped copy contributes none.
This lets JAX return the required zero at the integer root while preserving the slope encoded by the surrounding expression.
It also normalizes a possible negative zero into positive zero.
Low-precision inputs are evaluated in float32
For bfloat16 and float16, the real implementation first converts the input to float32.
The range reduction and polynomial are evaluated there, and the result is converted back to the original type.
This avoids asking a low-precision format to carry the intermediate coefficients and remainder arithmetic.
The change does not promise that every real dtype has identical output bits. It gives low-precision inputs a more accurate working format before final rounding.
Complex inputs keep a separate algorithm
The additive integer-plus-remainder construction is designed for real inputs.
The patch extracts the existing complex strategy into _sinc_complex().
Complex values still use:
direct sin(pi × x) / (pi × x)outside a small disk and a Taylor series near the origin.
That separation matters when interpreting the scope of the PR. The new full-range additive reduction and the published real-valued ULP bounds do not automatically establish the same behavior for complex inputs.
The tests cover roots, full range, and derivatives
The numerical reference test builds 10,000 random float64 values across the floating-point range and adds explicit corner cases:
positive and negative zero
positive and negative infinity
NaN
positive and negative integers
the adjacent representable values on both sides of those integersThe independent NumPy reference also uses additive range reduction. It is compared with mpmath.sincpi, and the test requires the reference to stay within three float64 ULPs for the sampled data.
The implementation accuracy test then checks jnp.sinc against that reference and mpmath across JAX’s supported real floating-point types and hardware-specific bounds.
The published thresholds include:
bfloat16 and float16:
0.5 ULP in the documented table
float32:
2.5 to 4.0 ULP depending on backend
float64:
2.0 ULP in the sampled CPU and GPU testsThe float64 table is empirical: sampling cannot prove the true worst-case input across the entire 64-bit domain.
The regression also differentiates four times
The higher-order regression evaluates 1st through 4th derivatives on 1,001 points across:
[-0.5, 0.5]It constructs an independent derivative reference from the Taylor series:
sum over k of
(-1)^k × pi^(2k)
× derivative factorial terms
× x^(2k-m)Each JAX derivative is compared with that reference using absolute and relative tolerances of 1e-4.
This directly addresses the long-standing issue that reported unstable fourth and higher derivatives near zero.
The test does not cover every derivative order. The implementation comments explain why the chosen polynomial degree is intended to retain accuracy through the tested orders and at zero through higher even derivatives.
A local arithmetic check separates the old and new boundaries
I reproduced the two algorithms in a local CPU JAX environment using the deleted real-valued branch from the PR diff and the merged float32 formula.
The check was not a build of the merged JAX commit.
For selected values, it produced:
x = 1.0
old branch:
-2.7827534e-08
new formula:
+0.0For the next float32 above one:
high-precision reference:
approximately -1.19209275e-07
old branch:
-1.79609486e-07
new formula:
-1.19209275e-07For the largest finite float32:
old branch:
NaN
new formula:
+0.0Across 1,001 float32 points in [-0.5, 0.5], the locally transcribed new formula also reduced the maximum absolute fourth-derivative error against the same Taylor reference compared with the deleted branch.
These are isolated arithmetic checks, not upstream test results or multi-backend validation.
What the patch does not establish
The pull request contains no performance benchmark.
The real path now executes rounding, parity detection, several polynomial operations, conditionals, and a division. Depending on backend and input distribution, that may have a different cost from a hardware sine call.
The available evidence supports improved numerical behavior. It does not support a universal speedup.
The tests also provide sampled ULP bounds rather than a formal exhaustive proof, especially for float64.
The public GitHub import contains no visible review discussion, so the internal review reasoning cannot be reconstructed from the PR thread.
The numerical contract restored by the change
A mathematical identity is not automatically a stable floating-point algorithm.
The corrected path is:
choose the nearest integer center
↓
subtract it exactly to obtain a small remainder
↓
evaluate a derivative-aware polynomial on that remainder
↓
restore parity and divide by the original input
↓
preserve exact roots, signed-zero behavior, and derivativesArgument reduction should preserve the structure that floating-point multiplication would erase.
JAX PR #40803 applies that rule to jnp.sinc, replacing a familiar formula with an implementation designed for roots, extreme inputs, and repeated automatic differentiation.
Sources
https://github.com/jax-ml/jax/pull/40803
https://github.com/jax-ml/jax/commit/bc6705044994785c356f7c39b068186647f120f7
https://github.com/jax-ml/jax/issues/5094
https://github.com/jax-ml/jax/blob/bc6705044994785c356f7c39b068186647f120f7/jax/_src/numpy/ufuncs.py
https://github.com/jax-ml/jax/blob/bc6705044994785c356f7c39b068186647f120f7/tests/lax_numpy_test.py
https://github.com/jax-ml/jax/blob/bc6705044994785c356f7c39b068186647f120f7/docs/accuracy.md
https://github.com/jax-ml/jax/blob/bc6705044994785c356f7c39b068186647f120f7/CHANGELOG.md