def_deriv rules now honor effective_order via scan path

_deriv_prop_propagate (used by every primitive registered through
def_deriv, currently just lax.erf_p) previously used
lax.conv_general_dilated, which is a single fused XLA op with no
runtime control flow — so a tracer-valued effective_order could not
short-circuit any work.

Switch path on the parameter:
  - effective_order is None: keep conv_general_dilated. Optimal for
    full-K computation, no overhead vs the prior implementation.
  - effective_order is not None: rewrite the same Faà di Bruno
    recurrence as a _jet_scan body, so iterations k > effective_order
    become no-ops via the existing lax.cond guard.

def_deriv now also calls _register_eff_order_rule(prim) so the
trace's process_primitive plumbs effective_order into the rule. No-op
behaviorally for any caller that never passes effective_order.

Tests: tests/test_effective_order.py covers every primitive in
_rules_with_effective_order. For each, it asserts that
jet(f, ..., effective_order=k)[:k] equals jet(f, ...)[:k] across
multiple values of k, both eager and under jit. Plus a composition
test (exp(sin(x))) and a full-K-equals-no-hint test for erf. 86 new
tests, all passing.

README: clarify the three-bucket taxonomy now that the def_deriv
bucket is fully covered. Move erf into the aware list; list the
remaining handful of unsupported primitives explicitly.

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

Author: thisiscam

README.md

@@ -238,20 +238,25 @@ combine with `jax.vmap`.
     applies elementwise along the order axis — `(a + b)[k] = a[k] + b[k]`
     — so it is already a single `O(K)` XLA op over the series array.
     There is no convolution loop to short-circuit, and `effective_order`
-    has no effect. (Output entries above `effective_order` still hold
+    has no effect. Output entries above `effective_order` still hold
     correct values for these primitives, because the linear op computes
-    every entry uniformly.)
-  - **Nonlinear without an `effective_order`-aware rule** (`sin`, `cos`,
-    `sinh`, `cosh`, `square`, `sqrt`, `erf`, …): the rule computes the
+    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
+    `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.
-  - **Nonlinear with an `effective_order`-aware rule** (`exp`, `expm1`,
-    `log`, `log1p`, `pow`, `logistic`, `tanh`, `erf_inv`, `div`): 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.
+    ignored on this path. These can be made aware with the same
+    `_jet_scan` plumbing as the others — see the source for the pattern.
 - Because output entries above `effective_order` are unspecified for the
-  third category, treat the whole `series_out[effective_order:]` slice as
+  second category, treat the whole `series_out[effective_order:]` slice as
   garbage even if upstream linear ops would otherwise leave it intact.
 - The win is largest when the dominant cost is in the convolution-style
   loops that the hint short-circuits. For shallow programs the overhead of

jet_array/__init__.py

@@ -502,36 +502,66 @@ def def_deriv(prim, deriv):
     Define the jet rule for a primitive in terms of its first derivative.
     """
     jet_rules[prim] = partial(deriv_prop, prim, deriv)
+    _register_eff_order_rule(prim)
 
 
-def deriv_prop(prim, deriv, primals_in, series_in):
+def deriv_prop(prim, deriv, primals_in, series_in, _jet_effective_order=None):
     (x,) = primals_in
     (series,) = series_in
     primal_out = prim.bind(x)
-    c0, cs = jet(deriv, primals_in, series_in)
-    tail = _deriv_prop_propagate(c0, cs, x, series)
+    c0, cs = jet(deriv, primals_in, series_in,
+                 effective_order=_jet_effective_order)
+    tail = _deriv_prop_propagate(c0, cs, x, series,
+                                 effective_order=_jet_effective_order)
     return primal_out, tail
 
 
-def _deriv_prop_propagate(c0, cs, x, series):
+def _deriv_prop_propagate(c0, cs, x, series, effective_order=None):
     c = _prepend_primal(c0, cs)
     u = _prepend_primal(x, series)
-    u_scaled = jnp.arange(u.shape[0]) * u
-    u_pad = jnp.pad(
-        u_scaled,
-        ((len(c) - 1, 0),) + ((0, 0),) * (u.ndim - 1),
-        mode="constant",
-        constant_values=0,
-    )
-    full = lax.conv_general_dilated(
-        u_pad[None, None, :],
-        c[::-1][None, None, :],
-        window_strides=(1,),
-        padding="VALID",
-        dimension_numbers=("NCW", "IOW", "NCW"),
-    )[0, 0]
-    v_tail = full[1 : len(u)] / jnp.arange(1, len(u))
-    return v_tail
+
+    if effective_order is None:
+        # Fast path: single fused conv, optimal for full-K computation.
+        u_scaled = jnp.arange(u.shape[0]) * u
+        u_pad = jnp.pad(
+            u_scaled,
+            ((len(c) - 1, 0),) + ((0, 0),) * (u.ndim - 1),
+            mode="constant",
+            constant_values=0,
+        )
+        full = lax.conv_general_dilated(
+            u_pad[None, None, :],
+            c[::-1][None, None, :],
+            window_strides=(1,),
+            padding="VALID",
+            dimension_numbers=("NCW", "IOW", "NCW"),
+        )[0, 0]
+        return full[1 : len(u)] / jnp.arange(1, len(u))
+
+    # Scan path: same Faà di Bruno recurrence
+    #   v[k] = (1/k) * sum_{j=1..k} j * u[j] * c[k-j]
+    # rewritten as a scan body so _jet_scan can short-circuit iterations
+    # k > effective_order.
+    n = u.shape[0]
+    if n == 1:
+        return jnp.zeros_like(series)
+
+    j_idx = jnp.arange(n, dtype=u.dtype).reshape((n,) + (1,) * (u.ndim - 1))
+    u_scaled = j_idx * u
+    pad_width = [(n - 2, 0)] + [(0, 0)] * (u.ndim - 1)
+    u_pad = jnp.pad(u_scaled, pad_width, mode="constant")[::-1]
+
+    v = jnp.zeros_like(u)
+
+    def body(v_acc, k):
+        u_slice = lax.dynamic_slice_in_dim(u_pad, n - 1 - k, n - 1, axis=0)
+        v_k = jnp.einsum("i...,i...->...", c[:-1], u_slice) / k
+        v_acc = v_acc.at[k].set(v_k)
+        return v_acc, None
+
+    v_final, _ = _jet_scan(body, v, jnp.arange(1, n, dtype=jnp.int32),
+                           effective_order)
+    return v_final[1:]
 
 
 def_deriv(

tests/test_effective_order.py

@@ -0,0 +1,153 @@
+"""Contract tests for ``effective_order``.
+
+For every primitive registered with ``_register_eff_order_rule``, verify
+that calling ``jet(f, primals, series, effective_order=k)`` produces
+output series whose first ``k`` entries match the full computation
+``jet(f, primals, series)``.
+
+Entries above index ``k`` are unspecified by contract (rules in the
+"nonlinear with aware rule" bucket leave them garbage), so this test
+checks only the ``[:k]`` slice.
+"""
+
+import pytest
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import lax
+
+from jet_array import jet
+import jet_array
+
+
+RTOL = 1e-10
+ATOL = 1e-12
+
+K = 12
+
+
+def _safe_x0(name):
+    """Pick a primal at which the function is well-defined and well-conditioned."""
+    return {
+        "log": 1.3,
+        "log1p": 0.5,
+        "div": 1.7,        # exercised via the f below
+        "pow": 1.4,
+        "erf_inv": 0.3,    # erf_inv defined on (-1, 1)
+        "tanh": 0.4,
+        "logistic": 0.4,
+        "exp": 0.3,
+        "expm1": 0.3,
+        "erf": 0.5,
+        "sin": 0.5,
+        "cos": 0.5,
+        "sinh": 0.4,
+        "cosh": 0.4,
+    }.get(name, 0.5)
+
+
+def _univariate_fn(name):
+    """Map primitive name to a univariate test function exercising that primitive.
+
+    For non-univariate primitives (mul, dot_general, conv_general_dilated)
+    we return None; those are exercised separately.
+    """
+    if name == "log":
+        return lambda x: lax.log(x)
+    if name == "log1p":
+        return lambda x: lax.log1p(x)
+    if name == "exp":
+        return lambda x: lax.exp(x)
+    if name == "expm1":
+        return lambda x: lax.expm1(x)
+    if name == "logistic":
+        return lambda x: lax.logistic(x)
+    if name == "tanh":
+        return lambda x: lax.tanh(x)
+    if name == "erf":
+        return lambda x: lax.erf(x)
+    if name == "erf_inv":
+        return lambda x: lax.erf_inv(x)
+    if name == "sin":
+        return lambda x: lax.sin(x)
+    if name == "cos":
+        return lambda x: lax.cos(x)
+    if name == "sinh":
+        return lambda x: lax.sinh(x)
+    if name == "cosh":
+        return lambda x: lax.cosh(x)
+    if name == "pow":
+        return lambda x: lax.pow(x, jnp.float64(1.7))
+    if name == "div":
+        return lambda x: 2.0 / (1.0 + x)
+    return None
+
+
+def _eff_order_aware_univariate_names():
+    return sorted(
+        p.name for p in jet_array._rules_with_effective_order
+        if _univariate_fn(p.name) is not None
+    )
+
+
+@pytest.mark.parametrize("name", _eff_order_aware_univariate_names())
+@pytest.mark.parametrize("eff", [1, 2, 4, 8, K])
+def test_effective_order_prefix_matches_full(name, eff):
+    """series[:eff] with effective_order=eff equals series[:eff] without it."""
+    fn = _univariate_fn(name)
+    x0 = jnp.asarray(_safe_x0(name), dtype=jnp.float64)
+    series_in = jnp.zeros(K, dtype=jnp.float64).at[0].set(1.0)
+
+    _, s_full = jet(fn, (x0,), (series_in,))
+    _, s_eff = jet(fn, (x0,), (series_in,), effective_order=jnp.array(eff))
+
+    np.testing.assert_allclose(
+        s_eff[:eff], s_full[:eff], rtol=RTOL, atol=ATOL,
+        err_msg=f"{name} eff={eff}: prefix mismatch",
+    )
+
+
+@pytest.mark.parametrize("name", _eff_order_aware_univariate_names())
+def test_effective_order_under_jit(name):
+    """effective_order works under jax.jit with a tracer-valued eff."""
+    fn = _univariate_fn(name)
+    x0 = jnp.asarray(_safe_x0(name), dtype=jnp.float64)
+    series_in = jnp.zeros(K, dtype=jnp.float64).at[0].set(1.0)
+
+    @jax.jit
+    def f(x, k):
+        return jet(fn, (x,), (series_in,), effective_order=k)
+
+    _, s_ref = jet(fn, (x0,), (series_in,))
+    for eff in [1, 4, K]:
+        _, s = f(x0, jnp.array(eff))
+        np.testing.assert_allclose(
+            s[:eff], s_ref[:eff], rtol=RTOL, atol=ATOL,
+            err_msg=f"{name} (jit) eff={eff}: prefix mismatch",
+        )
+
+
+def test_effective_order_in_composition():
+    """A composed function (sin then exp) propagates effective_order correctly."""
+    fn = lambda x: lax.exp(lax.sin(x))
+    x0 = jnp.asarray(0.5, dtype=jnp.float64)
+    series_in = jnp.zeros(K, dtype=jnp.float64).at[0].set(1.0)
+
+    _, s_ref = jet(fn, (x0,), (series_in,))
+    for eff in [1, 3, 6, K]:
+        _, s = jet(fn, (x0,), (series_in,), effective_order=jnp.array(eff))
+        np.testing.assert_allclose(
+            s[:eff], s_ref[:eff], rtol=RTOL, atol=ATOL,
+            err_msg=f"compose eff={eff}: prefix mismatch",
+        )
+
+
+def test_effective_order_full_K_equals_no_hint():
+    """effective_order=K (the full size) must produce results bit-equal to no hint."""
+    fn = lambda x: lax.erf(x)
+    x0 = jnp.asarray(0.4, dtype=jnp.float64)
+    series_in = jnp.zeros(K, dtype=jnp.float64).at[0].set(1.0)
+
+    _, s_no = jet(fn, (x0,), (series_in,))
+    _, s_K = jet(fn, (x0,), (series_in,), effective_order=jnp.array(K))
+    np.testing.assert_allclose(s_no, s_K, rtol=1e-12, atol=1e-14)