Skip to content

Add derivative_contractor via gradient_code SSA path - #81

Closed
dpsanders wants to merge 4 commits into
add-ssa-functionfrom
add-derivative-contractor
Closed

dpsanders wants to merge 4 commits into
add-ssa-functionfrom
add-derivative-contractor

Conversation

@dpsanders

Copy link
Copy Markdown
Member

Summary

  • Adds derivative_ssa(ex, diff_var, vars) — builds an SSAFunction for the partial derivative ∂ex/∂diff_var using reverse-mode AD (gradient_code)
  • Adds derivative_contractor(ex, diff_var, vars) — builds a compiled forward-backward contractor for the derivative
  • Handles gradient code's non-elementary assignments by binarizing the SSA (decomposing n-ary and compound-argument operations into unary/binary ops via CSE)
  • Filters non-reversible assignments (constant seeds, identity copies from gradient initialization) from the reverse pass

This builds on #80 (SSAFunction type). The gradient SSA is binarized and then fed through forward_backward_contractor via the SSAFunction interface.

Example

@variables x y
ex = x^2 + y^2

# Build contractor for ∂(x²+y²)/∂x = 2x
C = derivative_contractor(ex, x, [x, y])

# Contract: 2x ∈ [0,0] → x = 0
C(IntervalBox(-10..10, 2), interval(0, 0))
# → (([0,0], [-10,10]), [-20,20])

Test plan

  • All 18 ReversePropagation tests pass (was 12, +6 new for derivative_contractor)
  • All 16 IntervalConstraintProgramming tests pass (downstream consumer)
  • Simple polynomial: d/dx(x²+y²) = 2x, equality and interval constraints
  • Transcendental: d/dx(sin(xy)+x²) produces contracted result
  • Derivative w.r.t. second variable works correctly

🤖 Generated with Claude Code

@dpsanders
dpsanders force-pushed the add-derivative-contractor branch from 8c8313a to 81a8ef6 Compare April 19, 2026 19:08
dpsanders and others added 4 commits April 19, 2026 14:29
Avoids world-age issues when building and invoking a generated gradient
or forward-backward contractor in the same dynamic extent.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Patch bump for the RuntimeGeneratedFunctions swap: no changes to the
exported API's calling convention, just a correctness fix (no more
world-age errors) and a substantial build-time speedup.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Build forward-backward contractors for partial derivatives via reverse-mode
AD. The gradient_code SSA is binarized (decomposed into elementary ops) and
then fed into forward_backward_contractor through the SSAFunction interface.

- Add derivative_ssa(ex, diff_var, vars): builds SSAFunction for ∂ex/∂var
- Add derivative_contractor(ex, diff_var, vars): builds compiled contractor
- Add _binarize_ssa: decomposes n-ary and compound-argument assignments
  from gradient_code into elementary binary/unary ops for rev()
- Filter non-call assignments (constant/identity) from reverse pass in
  forward_backward_code, since they can't be reversed
- Export derivative_ssa and derivative_contractor

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Five independent fixes plus a registry-wide regression test:

- Drop `sign`, `max`, `min` from `unary_functions`. `sign_rev` is not
  defined anywhere (the registration has always been broken), and
  `max`/`min` are binary, not unary — ChainRules' adjoint for them
  synthesises boolean branch selectors (`!(y<x)`) which aren't
  reversible ops.
- Override ChainRules' `abs` adjoint with `x/abs(x)`; the default uses
  `sign(x)`, whose `_rev_sign` we just removed.
- Override ChainRules' `exp2`/`exp10`/`log2`/`log10` adjoints to use
  plain `Float64` logs instead of the `IrrationalConstants.Logtwo`/
  `Logten` that `IntervalArithmetic` refuses to accept.
- Extend `IntervalArithmetic.intersect_interval` with `(Real, Interval)`
  and `(Real, BareInterval)` methods so the constraint intersect works
  when the output variable comes through as a plain `Real` (e.g. from
  `d/dx(x^exact(2))`).
- Promote integer seeds (`_ā := 1`, unassigned `:= 0`) in the binarized
  derivative SSA to intervals. Without this, identity-propagating
  gradients such as `d/dx(x-y) = 1` leave `Int`-valued entries that
  drive `IntervalContractors`' `mul_rev(a,b,c) = mul_rev(promote(a,b,c)...)`
  fallback into infinite recursion on `(Int, Int, Int)`.

Also fixes PR 81's stale destructuring of `gradient_code` (it now
returns an `SSAFunction` rather than a tuple, post-#80).

Adds a regression test that loops over every entry in the binary and
unary registries and builds+runs a `derivative_contractor` — catches
each of these regressions.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@dpsanders
dpsanders force-pushed the add-derivative-contractor branch from 81a8ef6 to 0e227af Compare April 19, 2026 19:30
@dpsanders

Copy link
Copy Markdown
Member Author

Closing in favour of a split: the actual derivative_contractor needs subgradient-aware interval reverses for abs/max/min (see #894 context and the ChainRules discussion on this PR), and that's cleaner as its own PR. The non-interval fixes (derivative hardening: trim the op registry, override scalar rules that produced non-reversible adjoints, expose binarize_ssa) will land in a separate smaller PR shortly.

@dpsanders dpsanders closed this Apr 19, 2026
@dpsanders
dpsanders deleted the add-derivative-contractor branch April 19, 2026 22:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant