@@ -2,18 +2,41 @@
 
 Array-form Taylor-mode automatic differentiation in JAX.
 
-`jet_array` propagates truncated Taylor polynomials through arbitrary JAX
-computations, returning all coefficients up to a specified order in a single
-forward pass. This is the higher-order analogue of `jax.jvp`: where `jvp`
-takes a primal and one tangent and returns one derivative, `jet` takes a
-primal and a series of tangents and returns the full Taylor expansion.
-
-This package is a derivative of JAX's experimental `jax.experimental.jet`
-([Bettencourt, Johnson, Duvenaud 2019][taylor-mode]) with the series stored as
-a single leading-axis array rather than a Python tuple. The array layout
-enables `jit`/`vmap`/`scan` over the order axis, and adds an
-`effective_order` parameter for dynamic computation depth at static shapes —
-useful when the truncation order varies across a batch.
+`jet_array` propagates a truncated Taylor polynomial of order `K` through an
+arbitrary JAX computation in a single forward pass, returning all `K` Taylor
+coefficients of the output as one leading-axis array. This is the higher-order
+analogue of `jax.jvp`: where `jvp` takes one tangent and returns one
+derivative, `jet` takes a series of `K` coefficients and returns the full
+order-`K` Taylor expansion.
+
+The package extends `jax.experimental.jet`
+([Bettencourt, Johnson, Duvenaud 2019][taylor-mode]) in two ways:
+
+1. **Array storage.** Series are stored as a single `jnp.ndarray` along axis
+   0 instead of a Python tuple. This makes the order axis a first-class JAX
+   dimension that you can `jit`, `vmap`, and `scan` over.
+2. **`effective_order` parameter.** A runtime hint that lets a single jitted
+   program execute with a dynamic truncation depth — useful when each
+   element of a `vmap`-batched call needs a different number of coefficients.
+
+## When is Taylor-mode AD useful?
+
+Taylor-mode (jet) gives you the first `K` derivatives of a scalar function
+in `O(K)` work. Computing the same `K` derivatives by repeatedly nesting
+`jax.grad` costs `O(K²)` (the depth of the trace doubles each time). So
+Taylor-mode is the right tool when:
+
+- You need many derivatives at one point — for example, `K=20` to evaluate a
+  Bell polynomial, an ODE Taylor integrator, or a high-order moment.
+- The program you are differentiating contains the same primitives many
+  times — Taylor-mode amortizes one trace, the nested-grad approach traces
+  `K` times.
+- You need the full Taylor *polynomial* (with factorial-divided
+  coefficients) rather than the unscaled derivatives — for example, to
+  evaluate `f(x + h)` as a series in `h`.
+
+For first or second derivatives, ordinary `jax.grad` / `jax.jacrev` is
+faster — Taylor-mode pays only when `K ≥ ~3`.
 
 ## Install
 
@@ -25,62 +48,252 @@ pip install jet-array
 `jax>=0.8,<0.9`. Compatibility across JAX versions will be tightened as the
 package matures.
 
-## Quickstart
+## API
+
+```python
+jet_array.jet(fun, primals, series, effective_order=None) -> (primal_out, series_out)
+```
+
+- **`fun`** — a JAX-traceable callable.
+- **`primals`** — a tuple of input primal values, one per positional argument
+  of `fun`. Each must be a leaf (scalar or array, not a pytree).
+- **`series`** — a tuple of arrays, one per primal. `series[i][k-1]` is the
+  `k`-th Taylor coefficient of the `i`-th input along the path you are
+  expanding around. The trailing dimensions of `series[i]` must match
+  `primals[i]`. The order `K` of the expansion is the leading axis length;
+  it must be the same for all primals.
+- **`effective_order`** — optional. See below.
+- **Returns**: `(primal_out, series_out)` where `primal_out = fun(*primals)`
+  and `series_out[k-1]` is the `k`-th Taylor coefficient of the output, in
+  the same convention as the input.
+
+### Coefficient convention
+
+`series[k-1]` is the **Taylor coefficient**: if `f` is being expanded along
+`x(t) = x₀ + s₁·t + s₂·t² + …`, then
+
+  `f(x(t)) = f(x₀) + Σ_{k≥1} series_out[k-1] · t^k`.
+
+When you set `series_in = (1, 0, 0, …, 0)` and a single primal, this reduces
+to expanding `f` directly around `x₀`, and the relationship to derivatives
+is
+
+  `series_out[k-1] = f^(k)(x₀) / k!`.
+
+To recover the unscaled `k`-th derivative, multiply by `k!`. (This differs
+from `jax.experimental.jet`, which returns derivative coefficients without
+the factorial scaling. The relationship is documented and tested in
+`tests/test_against_jax_jet.py`.)
+
+## Examples
+
+### Univariate: high-order derivatives at one point
 
 ```python
+import math
+import jax
 import jax.numpy as jnp
 from jet_array import jet
 
+jax.config.update("jax_enable_x64", True)
+
 def f(x):
     return jnp.exp(jnp.sin(x))
 
-# 5th-order Taylor series of f at x=0.5, with input series [1, 0, 0, 0, 0]
-# (i.e. expand around x = 0.5 + t).
-x0 = 0.5
+x0 = jnp.asarray(0.5)
+K = 8                                        # expansion order
+series_in = jnp.zeros(K).at[0].set(1.0)      # direction: x(t) = x0 + t
+
+primal, series = jet(f, (x0,), (series_in,))
+
+# series[k-1] = f^(k)(x0) / k!
+# Multiply by k! to get the unscaled k-th derivative.
+for k in range(1, K + 1):
+    print(f"f^({k})(x0) = {float(series[k-1]) * math.factorial(k):.6f}")
+```
+
+### Bivariate: directional Taylor expansion
+
+```python
+def g(x, y):
+    return jnp.exp(x * y)
+
+x0 = jnp.asarray(0.5)
+y0 = jnp.asarray(0.3)
+
+# Expand along the path (x0 + t, y0): output series gives ∂^k g/∂x^k / k!
+sx = jnp.zeros(4).at[0].set(1.0)
+sy = jnp.zeros(4)                            # y stays constant
+primal, series = jet(g, (x0, y0), (sx, sy))
+# series[0] = ∂g/∂x        = y₀ · exp(x₀ y₀)
+# series[1] = ∂²g/∂x² / 2! = y₀² / 2 · exp(x₀ y₀)
+
+# Diagonal direction (x0+t, y0+t): mixed derivatives appear.
+sy_diag = jnp.zeros(4).at[0].set(1.0)
+_, series_diag = jet(g, (x0, y0), (sx, sy_diag))
+# series_diag[k-1] = (1/k!) · sum over multi-indices |α|=k of  ∂^α g · α-coefficients
+```
+
+### Inside `jax.jit`
+
+`jet` is fully traceable. The order `K` is part of the input shape, so a
+single jit-compiled program handles any computation at that order:
+
+```python
+@jax.jit
+def taylor(x0):
+    series_in = jnp.zeros(8).at[0].set(1.0)
+    return jet(f, (x0,), (series_in,))
+
+primal, series = taylor(jnp.asarray(0.5))
+```
+
+If you call `taylor` with a different `K`, JAX retraces (because the input
+shape changed). To avoid that — see `effective_order` below.
+
+### Inside `jax.vmap`
+
+```python
+xs = jnp.linspace(0.0, 1.0, 100)
 series_in = jnp.zeros(5).at[0].set(1.0)
+primals, series = jax.vmap(
+    lambda x: jet(f, (x,), (series_in,))
+)(xs)
+# primals.shape == (100,), series.shape == (100, 5)
+```
+
+## `effective_order`: dynamic truncation under jit
+
+`effective_order` is an integer JAX scalar that tells `jet` how many
+coefficients of the output you actually intend to use. The output array
+still has shape `(K, ...)` — same as without the parameter — but `jet`
+skips work computing high-order entries inside primitives whose Taylor
+rules support the hint (currently `exp`, `expm1`, `log`, `log1p`, `pow`,
+`logistic`, `tanh`, `erf_inv`, `div`).
+
+The value of entries beyond `effective_order` is unspecified; treat them
+as garbage. The point is to *avoid recompilation* when the order varies
+across calls or across a `vmap`-batched dimension, since the array shape
+is fixed.
+
+### When you want this
+
+A motivating workload: nested-Archimedean copula likelihoods with
+per-observation censoring. Each observation needs a Bell-polynomial
+expansion whose effective order equals its number of *uncensored* leaves,
+which varies across the batch. Without `effective_order`, you face two
+bad options:
+
+- Pad to the maximum possible order for every observation — wasted work
+  on always-zero coefficients.
+- Re-trace and re-compile per observation — fatal under high cardinality.
+
+`effective_order` lets you compile once at `K = max possible order` and
+pay per-observation work proportional to that observation's actual order.
+
+### Example
 
-primal_out, series_out = jet(f, (x0,), (series_in,))
-# primal_out  = f(0.5)
-# series_out  = [f'(0.5), f''(0.5)/2!, f'''(0.5)/3!, f^(4)(0.5)/4!, f^(5)(0.5)/5!]
+```python
+import jax
+import jax.numpy as jnp
+from jet_array import jet
+
+K_MAX = 16                                    # static array size
+
+def slow_chain(x):
+    # A composition that uses several primitives whose rules respect
+    # effective_order — exp/log/logistic are good candidates.
+    y = jnp.exp(x)
+    y = jnp.log1p(y)
+    y = jnp.tanh(y)
+    return y
+
+@jax.jit
+def taylor(x, k_dyn):
+    series_in = jnp.zeros(K_MAX).at[0].set(1.0)
+    return jet(slow_chain, (x,), (series_in,), effective_order=k_dyn)
+
+# Both calls use the same compiled XLA program.
+# The second runs more work — it computes 12 coefficients instead of 4.
+p1, s1 = taylor(jnp.asarray(0.5), jnp.array(4))
+p2, s2 = taylor(jnp.asarray(0.5), jnp.array(12))
+# s1[:4] is the answer you can use; s1[4:] is unspecified.
+# s2[:12] is the answer; s2[12:] is unspecified.
 ```
 
+For workloads where the dynamic order varies per element of a
+`vmap`-batched dimension, pass an array of `effective_order` values and
+combine with `jax.vmap`.
+
+### Caveats
+
+- Only primitives in `_rules_with_effective_order` honor the hint. Other
+  primitives compute the full series. So for a function like `sin(x)`
+  alone, `effective_order` has no effect — `sin` always computes all `K`
+  coefficients.
+- Output entries with index `≥ effective_order` are not guaranteed to
+  contain anything specific (zero, garbage, or a partial answer). Slice
+  to `series_out[:effective_order]` before using them.
+- 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
+  the conditional may exceed the savings.
+
 ## Why arrays instead of tuples
 
-`jax.experimental.jet` stores the series as a Python tuple, which prevents
-`jit`/`vmap` from treating the order axis as a regular array dimension. This
-matters when:
+`jax.experimental.jet` stores the series as a Python tuple. That works for
+small `K` known at trace time but causes friction otherwise:
+
+- The order axis is invisible to JAX — you cannot `vmap` or `scan` over it.
+- A tuple of length `K` produces `K` separate jaxpr equations per
+  primitive, so trace time grows linearly in `K` even when the
+  computation is constant per coefficient.
+- `effective_order` makes no sense for a Python tuple.
+
+Storing series along axis 0 of a single array makes all of these go away.
+The cost is some boilerplate inside primitive rules to convolve along that
+axis, which `jet_array` provides.
+
+## Coverage
+
+Custom Taylor rules are provided for the JAX primitives that arise in
+typical scientific code: arithmetic and broadcasting, `exp`, `expm1`,
+`log`, `log1p`, `sin`, `cos`, `sinh`, `cosh`, `tanh`, `logistic`, `erf`,
+`erf_inv`, `pow`, `square`, `sqrt`, `div`, `dynamic_slice`,
+`dynamic_update_slice`, and `cumsum`/`cumprod`. Linear and zero-derivative
+primitives are handled generically. For untraced primitives the rule
+falls back to the standard convolution propagator.
+
+The full correctness suite checks every supported primitive against
+`jax.experimental.jet` at multiple expansion points and orders up to 20.
 
-- The series order is large (hundreds of coefficients) — Python-level
-  iteration becomes a tracing-time cost.
-- The order is data-dependent — tuples force re-tracing every time the
-  length changes.
-- You want to `vmap` or `scan` over a batch where each element has its own
-  effective order.
+## Limitations
 
-`jet_array` stores all coefficients along axis 0 of a single array,
-addressing all three.
+- Pytree primal inputs are not supported (each primal must be a leaf).
+- The package depends on `jax._src` internals and currently targets
+  JAX 0.8.x.
+- `effective_order` is a hint, not a guarantee — see the caveats above.
 
 ## Citation
 
-If you use `jet-array` in academic work, please cite:
+If you use `jet-array` in academic work, please cite the repository:
 
 ```bibtex
-@misc{yang2026copulaad,
-  title={Archimedean Copula Inference via Taylor-Mode AD},
-  author={Yang, Cambridge and Li, Dongdong},
-  year={2026},
-  note={arXiv preprint},
+@software{jet_array,
+  title  = {jet-array: array-form Taylor-mode automatic differentiation in JAX},
+  author = {Yang, Cambridge and Li, Dongdong},
+  year   = {2026},
+  url    = {https://github.com/thisiscam/jet-array},
 }
 ```
 
 The underlying Taylor-mode algorithm is from:
 
 ```bibtex
 @inproceedings{bettencourt2019taylor,
-  title={Taylor-mode automatic differentiation for higher-order derivatives in JAX},
-  author={Bettencourt, Jesse and Johnson, Matthew J. and Duvenaud, David},
-  booktitle={NeurIPS Program Transformations Workshop},
-  year={2019},
+  title     = {Taylor-mode automatic differentiation for higher-order derivatives in JAX},
+  author    = {Bettencourt, Jesse and Johnson, Matthew J. and Duvenaud, David},
+  booktitle = {NeurIPS Program Transformations Workshop},
+  year      = {2019},
 }
 ```
 
 
@@ -0,0 +1,51 @@
+import jax
+jax.config.update("jax_enable_x64", True)
+
+import pytest
+from jax import lax
+
+
+@pytest.fixture
+def factorial():
+    def _factorial(n):
+        if n <= 1:
+            return 1
+        return n * _factorial(n - 1)
+    return _factorial
+
+
+# Common univariate test functions: smooth, well-defined Taylor expansions.
+TEST_FUNCTION_PARAMS = [
+    ('erf', lambda x: lax.erf(x)),
+    ('exp', lambda x: lax.exp(x)),
+    ('log1p', lambda x: lax.log1p(x)),
+    ('expm1', lambda x: lax.expm1(x)),
+    ('sin', lambda x: lax.sin(x)),
+    ('cos', lambda x: lax.cos(x)),
+    ('sinh', lambda x: lax.sinh(x)),
+    ('cosh', lambda x: lax.cosh(x)),
+    ('tanh', lambda x: lax.tanh(x)),
+    ('logistic', lambda x: lax.logistic(x)),
+    ('square', lambda x: x**2),
+    ('compose_exp_sin', lambda x: lax.exp(lax.sin(x))),
+    ('compose_log_cosh', lambda x: lax.log(lax.cosh(x))),
+]
+
+
+MULTIVARIATE_TEST_PARAMS = [
+    ('bivariate_poly', lambda x, y: x**2 + x * y + y**2),
+    ('distance', lambda x, y: lax.sqrt(x**2 + y**2)),
+    ('exp_product', lambda x, y: lax.exp(x) * lax.exp(y)),
+    ('log_sum_exp', lambda x, y: lax.log(lax.exp(x) + lax.exp(y))),
+]
+
+
+EDGE_CASE_PARAMS = [
+    ('zero', 0.0),
+    ('small_positive', 1e-6),
+    ('small_negative', -1e-6),
+    ('mid', 0.5),
+    ('moderate', 1.5),
+    ('large_positive', 5.0),
+    ('large_negative', -5.0),
+]