Skip to content

Preserve despecialized parameters in adjoint VJPs - #1601

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:agent/autodespecialize-ad-parameters
Draft

Preserve despecialized parameters in adjoint VJPs#1601
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:agent/autodespecialize-ad-parameters

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member

What changed

Restore the parameter wrapper at SciMLSensitivity's central adjoint VJP entry points when the originating problem uses DespecializedParameters. AD backends may transform the concrete parameter value before it reaches vecjacobian!/vecjacobian; rebuilding the public, idempotent wrapper keeps the adjoint caches and backend arguments type-consistent.

This is required by ModelingToolkit's AutoDespecialize support in SciML/ModelingToolkit.jl#4919. The complementary transformed-call barrier fix is SciML/OrdinaryDiffEq.jl#4290.

Ignore this PR until reviewed by @ChrisRackauckas.

Failing before / passing after

The regression differentiates a plain ODEProblem{true, AutoDespecialize} using GaussAdjoint(EnzymeVJP()).

Before:

MethodError: no method matching EnzymeCore.Duplicated(::Vector{Float64}, ::SciMLBase.DespecializedParameters)

After, using registered DiffEqBase 7.15.0:

autodespecialize_adjoint_registered_diffeqbase=pass
grad=[-0.6065306597106055]

The gradient agrees with the analytic value -exp(-0.5).

Verification

GROUP=QA julia +1.12 --startup-file=no --project=. -e 'using Pkg; Pkg.test()'
Quality Assurance | 20 / 20 | 2m06.8s
Testing SciMLSensitivity tests passed

The exact SciMLSensitivity Core8 downstream environment from SciML/ModelingToolkit.jl#4919 was run with this patch and SciML/OrdinaryDiffEq.jl#4290:

before: Core 8 | 35 pass | 5 error | 5 broken
after:  Core 8 | 45 pass | 5 broken | 50 total | 108m24.9s
Testing SciMLSensitivity tests passed

Runic 1.8, added-line typos, and git diff --check pass.

Not verified

The full Core8 group was run with both owner fixes together; the new focused regression was additionally run with registered DiffEqBase 7.15.0 and passed independently. GPU and allowed-to-fail prerelease jobs were not run locally. No public API or documentation changed, so no docs build was required.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Copy link
Copy Markdown
Member Author

A fresh SciMLBase integration run on current dependency heads reproduces the same SciMLSensitivity Core8 regression:

  • Core 8 | 35 pass | 5 error | 5 broken
  • repeated a parameter-despecialization barrier requires exactly one parameter wrapper
  • adjoint parameter conversions involving DespecializedParameters

This is independent current-CI evidence for the paired owner fixes SciML/OrdinaryDiffEq.jl#4290 and this PR.

CI job: https://github.com/SciML/SciMLBase.jl/actions/runs/32365988822/job/96415492868

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.

2 participants