Commit 6319a06
committed
Replace array jacobian() with jvp()/vjp() forward & reverse modes
Drop the jacobian(expr, [cols]) -> List<Float64> form: a nested array column
breaks the long/tidy data model (a cell should be one value aligned to its
coordinates). The same Jacobian is expressed in-model as several scalar
columns, e.g. grad(f, x) AS dfdx, grad(f, y) AS dfdy.
Add forward- and reverse-mode gradients as scalar SQL functions:
* jvp(expr, column, tangent) -> d(expr)/d(column) * tangent (forward)
* vjp(expr, column, cotangent) -> cotangent * d(expr)/d(column) (reverse)
A multi-input directional derivative is the sum of per-input jvp terms; both
stay scalar, so they round-trip cleanly through Substrait and back to xarray.
Engine: unify grad and jvp behind a single `linearize` (forward-mode chain
rule with a pluggable leaf rule) — grad is a one-hot seed, jvp an arbitrary
seed per input. This mirrors JAX's structure and removes rule duplication.
vjp is cotangent * grad; for a scalar output forward and reverse coincide
(asserted by a jvp/vjp agreement test), differing only in seed placement.
Tests: 15 Rust unit tests and 11 Python integration tests (incl. jvp/vjp
semantics, the multi-input sum, and jvp==vjp for a unit seed), all checked
against numpy analytic derivatives. fmt/clippy/ruff/mypy clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017mDoFJgsm9kS7SicGoCVF61 parent 724099a commit 6319a06
4 files changed
Lines changed: 224 additions & 186 deletions
0 commit comments