Bucket C: plumb effective_order through atan2, integer_pow, cumprod/max/min

These rules call jet() or _deriv_prop_propagate() internally and
otherwise just forward to already-aware primitives. Pure plumbing:

- _atan2_taylor: forward to two inner jet() calls and to
  _deriv_prop_propagate
- _integer_pow_taylor: forward to inner jet() (both branches)
- _cumulative_jet_rule: forward to inner jet() over associative_scan
  (registered for cumprod, cummax, cummin)

Each gets _register_eff_order_rule(prim).

Also classify the remaining rules. The fan-out rules (abs, max, min,
gather, select_n, scatter-add, reduce_max, reduce_min) all use
jax.vmap over the order axis — structurally linear, like add/sub.
They don't need the effective_order plumbing because per-entry cost
is O(1) and the contract permits garbage above the threshold. README
updated to reflect both the expanded aware list (now 22 primitives)
and this third bucket.

Tests: extend tests/test_effective_order.py to cover all 5 new
primitives (atan2, integer_pow, cumprod, cummax, cummin). 30 new
parametric cases, all passing. Total suite: 638 passed, 5 skipped.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: thisiscam

README.md

@@ -242,19 +242,21 @@ combine with `jax.vmap`.
     correct values for these primitives, because the linear op computes
     every entry uniformly.
   - **Nonlinear with an `effective_order`-aware rule.** Currently
-    `exp`, `expm1`, `log`, `log1p`, `pow`, `logistic`, `tanh`, `erf_inv`,
-    `div`, `sin`, `cos`, `sinh`, `cosh`, `mul`, `dot_general`,
-    `conv_general_dilated`, `erf`, and any future `def_deriv`-registered
-    primitive. The convolution scan short-circuits iterations beyond
+    `exp`, `expm1`, `log`, `log1p`, `pow`, `integer_pow`, `logistic`,
+    `tanh`, `erf_inv`, `div`, `sin`, `cos`, `sinh`, `cosh`, `mul`,
+    `dot_general`, `conv_general_dilated`, `erf`, `atan2`, `cumprod`,
+    `cummax`, `cummin`, and any future `def_deriv`-registered primitive.
+    The convolution scan short-circuits iterations beyond
     `effective_order`. Output entries above `effective_order` are
     unspecified — slice to `series_out[:effective_order]` before using
     them.
-  - **Nonlinear without an `effective_order`-aware rule.** A few rules
-    (`abs`, `max`, `min`, `atan2`, `integer_pow`, `cumprod`, `cummax`,
-    `cummin`, the `reduce_*` family, `scatter-add`) currently compute the
-    full `O(K²)` convolution regardless. `effective_order` is silently
-    ignored on this path. These can be made aware with the same
-    `_jet_scan` plumbing as the others — see the source for the pattern.
+  - **Linear-style fan-out rules** (`abs`, `max`, `min`, `gather`,
+    `select_n`, `scatter-add`, `reduce_max`, `reduce_min`): each rule
+    applies an op independently to every series index via `vmap`, the
+    same shape as `add`/`sub` in the linear bucket. They don't
+    short-circuit, but they don't need to — the per-entry cost is O(1)
+    and the contract is satisfied because garbage-in-garbage-out for
+    indices above `effective_order` is fine.
 - Because output entries above `effective_order` are unspecified for the
   second category, treat the whole `series_out[effective_order:]` slice as
   garbage even if upstream linear ops would otherwise leave it intact.

jet_array/__init__.py

@@ -479,7 +479,8 @@ def _dynamic_update_slice_jet_rule(primals_in, series_in, **params):
 
 
 def _cumulative_jet_rule(
-    primals_in, series_in, *, axis: int, reverse: bool, combine_fn: Callable
+    primals_in, series_in, *, axis: int, reverse: bool, combine_fn: Callable,
+    _jet_effective_order=None,
 ):
     # Irrespective of backend, we always use the parallel prefix scan
     # implementation when differentiating because reduce_window is not
@@ -488,13 +489,16 @@ def _cumulative_jet_rule(
         partial(lax.associative_scan, combine_fn, axis=axis, reverse=reverse),
         primals_in,
         series_in,
+        effective_order=_jet_effective_order,
     )
 
 
 deflinear(lax.cumsum_p)
 jet_rules[lax.cumprod_p] = partial(_cumulative_jet_rule, combine_fn=lax.mul)
 jet_rules[lax.cummax_p] = partial(_cumulative_jet_rule, combine_fn=lax.max)
 jet_rules[lax.cummin_p] = partial(_cumulative_jet_rule, combine_fn=lax.min)
+for _p in (lax.cumprod_p, lax.cummax_p, lax.cummin_p):
+    _register_eff_order_rule(_p)
 
 
 def def_deriv(prim, deriv):
@@ -799,14 +803,17 @@ def _pow_by_squaring(x, n):
         return x * _pow_by_squaring(x * x, (n - 1) / 2)
 
 
-def _integer_pow_taylor(primals_in, series_in, *, y):
+def _integer_pow_taylor(primals_in, series_in, *, y, _jet_effective_order=None):
     if y == 0:
-        return jet(jnp.ones_like, primals_in, series_in)
+        return jet(jnp.ones_like, primals_in, series_in,
+                   effective_order=_jet_effective_order)
     else:
-        return jet(lambda x: _pow_by_squaring(x, y), primals_in, series_in)
+        return jet(lambda x: _pow_by_squaring(x, y), primals_in, series_in,
+                   effective_order=_jet_effective_order)
 
 
 jet_rules[lax.integer_pow_p] = _integer_pow_taylor
+_register_eff_order_rule(lax.integer_pow_p)
 
 
 def _logistic_taylor(primals_in, series_in, _jet_effective_order=None, **_):
@@ -924,18 +931,22 @@ def body_fun(v_acc, k):
 _register_eff_order_rule(lax.log_p)
 
 
-def _atan2_taylor(primals_in, series_in):
+def _atan2_taylor(primals_in, series_in, _jet_effective_order=None):
     x, y = primals_in
     primal_out = lax.atan2(x, y)
 
-    x, series = jet(lax.div, primals_in, series_in)
+    x, series = jet(lax.div, primals_in, series_in,
+                    effective_order=_jet_effective_order)
     one = lax_internal._const(x, 1)
-    c0, cs = jet(lambda x: lax.div(one, 1 + lax.square(x)), (x,), (series,))
-    tail = _deriv_prop_propagate(c0, cs, x, series)
+    c0, cs = jet(lambda x: lax.div(one, 1 + lax.square(x)), (x,), (series,),
+                 effective_order=_jet_effective_order)
+    tail = _deriv_prop_propagate(c0, cs, x, series,
+                                 effective_order=_jet_effective_order)
     return primal_out, tail
 
 
 jet_rules[lax.atan2_p] = _atan2_taylor
+_register_eff_order_rule(lax.atan2_p)
 
 
 def _div_taylor_rule(primals_in, series_in, _jet_effective_order=None, **_):

tests/test_effective_order.py

@@ -43,6 +43,11 @@ def _safe_x0(name):
         "cos": 0.5,
         "sinh": 0.4,
         "cosh": 0.4,
+        "atan2": 0.5,
+        "integer_pow": 0.7,
+        "cumprod": 0.6,
+        "cummax": 0.4,
+        "cummin": 0.4,
     }.get(name, 0.5)
 
 
@@ -80,6 +85,16 @@ def _univariate_fn(name):
         return lambda x: lax.pow(x, jnp.float64(1.7))
     if name == "div":
         return lambda x: 2.0 / (1.0 + x)
+    if name == "atan2":
+        return lambda x: lax.atan2(x, jnp.float64(2.0))
+    if name == "integer_pow":
+        return lambda x: x ** 3
+    if name == "cumprod":
+        return lambda x: lax.cumprod(jnp.stack([x, x * 2.0, x * 3.0]))
+    if name == "cummax":
+        return lambda x: lax.cummax(jnp.stack([x * 1.5, x * 0.5, x * 2.0]))
+    if name == "cummin":
+        return lambda x: lax.cummin(jnp.stack([x * 1.5, x * 0.5, x * 2.0]))
     return None