Skip to content

Commit bc36df3

Browse files
thisiscamclaude
andcommitted
README + tests: align with jax.experimental.jet's factorial_scaled flag
Recent jax.experimental.jet (PR jax-ml/jax#29676, merged 2025-07) accepts factorial_scaled=False, which returns the same Taylor coefficients as jet_array. Update the convention discussion in README accordingly, and switch the comparison tests to use factorial_scaled=False directly instead of post-hoc dividing standard-jet output by k!. The tests now assert direct coefficient equality rather than equality-after-rescaling. Also clarify the effective_order caveats. The hint applies to three primitive categories with different behavior: - linear / zero-derivative (add, sub, neg, broadcast, …): no effect, already O(K) elementwise - nonlinear without aware rule (sin, cos, sqrt, erf, …): silently ignored - nonlinear with aware rule (exp, log1p, pow, tanh, …): convolution scan short-circuits, entries above effective_order unspecified Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 56375ba commit bc36df3

2 files changed

Lines changed: 45 additions & 36 deletions

File tree

‎README.md‎

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -80,10 +80,15 @@ is
8080

8181
`series_out[k-1] = f^(k)(x₀) / k!`.
8282

83-
To recover the unscaled `k`-th derivative, multiply by `k!`. (This differs
84-
from `jax.experimental.jet`, which returns derivative coefficients without
85-
the factorial scaling. The relationship is documented and tested in
86-
`tests/test_against_jax_jet.py`.)
83+
To recover the unscaled `k`-th derivative, multiply by `k!`. This is the
84+
convention used in *Evaluating Derivatives* (Griewank & Walther, §13).
85+
86+
`jax.experimental.jet` exposes both conventions through its
87+
`factorial_scaled` keyword argument (defaults to `True`, returning
88+
derivative coefficients `f^(k)(x₀)`). Calling
89+
`jax.experimental.jet.jet(..., factorial_scaled=False)` returns the same
90+
Taylor coefficients as `jet_array`. The equivalence is tested
91+
coefficient-by-coefficient in `tests/test_against_jax_jet.py`.
8792

8893
## Examples
8994

@@ -227,13 +232,27 @@ combine with `jax.vmap`.
227232

228233
### Caveats
229234

230-
- Only primitives in `_rules_with_effective_order` honor the hint. Other
231-
primitives compute the full series. So for a function like `sin(x)`
232-
alone, `effective_order` has no effect — `sin` always computes all `K`
233-
coefficients.
234-
- Output entries with index `≥ effective_order` are not guaranteed to
235-
contain anything specific (zero, garbage, or a partial answer). Slice
236-
to `series_out[:effective_order]` before using them.
235+
- The primitives a JAX program uses fall into three categories:
236+
- **Linear / zero-derivative** (`add`, `sub`, multiplication by a
237+
constant, `neg`, broadcasting, slicing, comparisons, …): the operation
238+
applies elementwise along the order axis — `(a + b)[k] = a[k] + b[k]`
239+
— so it is already a single `O(K)` XLA op over the series array.
240+
There is no convolution loop to short-circuit, and `effective_order`
241+
has no effect. (Output entries above `effective_order` still hold
242+
correct values for these primitives, because the linear op computes
243+
every entry uniformly.)
244+
- **Nonlinear without an `effective_order`-aware rule** (`sin`, `cos`,
245+
`sinh`, `cosh`, `square`, `sqrt`, `erf`, …): the rule computes the
246+
full `O(K²)` convolution regardless. `effective_order` is silently
247+
ignored on this path.
248+
- **Nonlinear with an `effective_order`-aware rule** (`exp`, `expm1`,
249+
`log`, `log1p`, `pow`, `logistic`, `tanh`, `erf_inv`, `div`): the
250+
convolution scan short-circuits iterations beyond `effective_order`.
251+
Output entries above `effective_order` are unspecified — slice to
252+
`series_out[:effective_order]` before using them.
253+
- Because output entries above `effective_order` are unspecified for the
254+
third category, treat the whole `series_out[effective_order:]` slice as
255+
garbage even if upstream linear ops would otherwise leave it intact.
237256
- The win is largest when the dominant cost is in the convolution-style
238257
loops that the hint short-circuits. For shallow programs the overhead of
239258
the conditional may exceed the savings.

‎tests/test_against_jax_jet.py‎

Lines changed: 15 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,12 @@
11
"""Correctness tests against ``jax.experimental.jet``.
22
3-
``jax.experimental.jet`` returns derivative coefficients
4-
``[f(x), f'(x), f''(x), f'''(x), ...]``; ``jet_array`` returns Taylor
5-
coefficients ``[f(x), f'(x), f''(x)/2!, f'''(x)/3!, ...]``. So the
6-
relationship is::
7-
8-
jet_array_series[k] == standard_jet_series[k] / (k+1)! for k = 0, 1, ...
9-
10-
(Note the index shift: ``standard_jet`` returns the primal as
11-
``series[0]`` only when the user requests it; in our framing
12-
``primal_out`` is returned separately and ``series_out[k]`` is the
13-
``(k+1)``-th coefficient.)
14-
15-
Each test scales the standard-jet output by 1/k! and compares against
16-
``jet_array`` to within tight tolerance.
3+
Both libraries return Taylor coefficients ``[f'(x)/1!, f''(x)/2!, ...]``
4+
when invoked with matching conventions. ``jet_array.jet`` always uses
5+
this convention; ``jax.experimental.jet.jet`` returns it when called
6+
with ``factorial_scaled=False`` (with default ``True`` it returns the
7+
unscaled derivative coefficients ``[f'(x), f''(x), ...]``). These tests
8+
exercise the matching path and assert coefficient-by-coefficient
9+
equivalence.
1710
"""
1811

1912
import math
@@ -35,26 +28,22 @@
3528
ATOL = 1e-5
3629

3730

38-
def _scale_to_taylor(series_in_deriv):
39-
"""Convert standard-jet derivative coefficients to Taylor coefficients."""
40-
return [coeff / math.factorial(k + 1) for k, coeff in enumerate(series_in_deriv)]
41-
42-
4331
def _compare(name, x0, fn, order=5):
4432
"""Compare jet_array.jet against jax.experimental.jet at one point."""
4533
x0_arr = jnp.asarray(x0, dtype=jnp.float64)
4634
series_in_arr = jnp.zeros(order, dtype=jnp.float64).at[0].set(1.0)
4735
p_arr, s_arr = jet(fn, (x0_arr,), (series_in_arr,))
4836

4937
series_in_std = [1.0] + [0.0] * (order - 1)
50-
p_std, s_std = standard_jet.jet(fn, (x0_arr,), (series_in_std,))
51-
s_std_taylor = _scale_to_taylor(s_std)
38+
p_std, s_std = standard_jet.jet(
39+
fn, (x0_arr,), (series_in_std,), factorial_scaled=False
40+
)
5241

5342
np.testing.assert_allclose(
5443
p_arr, p_std, rtol=RTOL, atol=ATOL,
5544
err_msg=f"{name}@{x0}: primal mismatch",
5645
)
57-
for k, (a, b) in enumerate(zip(s_arr, s_std_taylor)):
46+
for k, (a, b) in enumerate(zip(s_arr, s_std)):
5847
np.testing.assert_allclose(
5948
a, b, rtol=RTOL, atol=ATOL,
6049
err_msg=f"{name}@{x0}: coefficient {k+1} mismatch",
@@ -92,11 +81,12 @@ def test_multivariate_matches_standard_jet(name, fn, order):
9281

9382
sx_std = [1.0] + [0.0] * (order - 1)
9483
sy_std = [1.0] + [0.0] * (order - 1)
95-
p_std, s_std = standard_jet.jet(fn, (x0, y0), (sx_std, sy_std))
96-
s_std_taylor = _scale_to_taylor(s_std)
84+
p_std, s_std = standard_jet.jet(
85+
fn, (x0, y0), (sx_std, sy_std), factorial_scaled=False
86+
)
9787

9888
np.testing.assert_allclose(p_arr, p_std, rtol=RTOL, atol=ATOL)
99-
for k, (a, b) in enumerate(zip(s_arr, s_std_taylor)):
89+
for k, (a, b) in enumerate(zip(s_arr, s_std)):
10090
np.testing.assert_allclose(
10191
a, b, rtol=RTOL, atol=ATOL,
10292
err_msg=f"{name}: coefficient {k+1} mismatch",

0 commit comments

Comments
 (0)