Repository navigation
Conversation
Composes the two AD passes introduced in the previous PRs: - One reverse-mode call (`gradient_code`) builds the SSA for the gradient. - `binarize_ssa` decomposes any n-ary adjoints the reverse pass left behind (e.g. `2 * _b̄ * x`) so `tangent_eq` can consume them. - `tangent_code` attaches a forward-mode tangent pass on top of the gradient SSA. - The compiled function invokes the tangent pass `n` times at runtime with unit-vector directions and assembles the `n × n` Hessian matrix. Uses exactly the same `unary_rule` / `binary_rule` table as the first two passes — a third source of truth would be a bug farm. Also fixes `tangent_eq` to handle non-call assignments (bare-variable identity copies, bare-constant seeds). The reverse-mode initialization (`_c̄ := 1`) and identity copies (`_b̄ := _c̄`) produced by `gradient_code` are non-call rhs's; the previous `tangent_eq` required every rhs to be a function call, which was fine for `cse_equations(ex)` output but blew up the moment a gradient SSA fed into forward-mode. Bumps to v0.8.0. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacked on #84 (forward-mode AD), which is stacked on #83 (derivative hardening). Adds
hessian(ex, vars)as the third AD entry point, reusing the same scalar rule table (unary_rule/binary_rule) asgradientandtangent.How it works
Forward-over-reverse:
gradient_code(ex, vars)— one reverse-mode pass. Yields an SSA whosevariables.gradient = [g₁, …, gₙ]are symbolic handles into the code.binarize_ssa— decompose any n-ary muls the adjoint left behind (e.g.2 * _b̄ * x) so the tangent pass can consume them.tangent_code— attach a forward-mode tangent pass on top. Produces tangent assignments for every SSA variable, including thegᵢs.(inputs, direction) → (value, gradient, gradient_tangent). Invoking it with directioneᵢgives rowiof the Hessian. The user-facinghessianclosure loops overi = 1..nand assembles then × nmatrix.API
Also in this PR
tangent_eqis hardened to accept non-call rhs's. Reverse-mode SSA contains_c̄ := 1(constant seed) and identity copies like_b̄ := _c̄; the previoustangent_eqassumed every rhs was a function call and errored the moment gradient SSA was used as forward-mode input. Tangent of a bare constant is 0, tangent of a bare variable is the variable's own tangent — both fall out ofdotteddispatching onRealvsBasicSymbolic.Test plan
x² + y²(diagonal[2, 2]).x·y(off-diagonal 1s).x²y + y³,sin(xy)— matches hand-derived Hessian.x² + y² + z² + xyz— dense 3×3.Bump to v0.8.0.
Performance note
The runtime cost is
ntangent-pass invocations perhessian(…)(inputs)call. Fine for the small-nsymbolic regime the package targets; denser cases can compile allntangent passes into a single SSA later.🤖 Generated with Claude Code