Skip to content

Fix Zygote gradient through scaled concretization#371

Open
AshtonSBradley wants to merge 5 commits intoSciML:masterfrom
AshtonSBradley:fix-zygote-concretize-gradient
Open

Fix Zygote gradient through scaled concretization#371
AshtonSBradley wants to merge 5 commits intoSciML:masterfrom
AshtonSBradley:fix-zygote-concretize-gradient

Conversation

@AshtonSBradley
Copy link
Copy Markdown
Contributor

@AshtonSBradley AshtonSBradley commented May 7, 2026

Summary

  • Fixes a Zygote sensitivity bug where update_coefficients(...) |> concretize for a scaled operator could produce a doubled gradient compared with the equivalent matrix expression.
  • Keeps ScaledOperator immutable, addressing the allocation concern from review.
  • Freezes updated ScalarOperator state into a private immutable scalar snapshot for the out-of-place scaled update path.
  • Adds AD semantic equivalence tests for scaled and added matrix operators across concretization, direct application, and mutating application.
  • Bumps the ArrayInterface lower bound so downgrade CI resolves a test-compatible Adapt version with ALLOW_RERESOLVE=false.

Fixes #305.

Explanation

Issue #305 reduces to a small operator expression of the form MatrixOperator(A1) + ScalarOperator(...) * MatrixOperator(A2). The primal value from update_coefficients(...) |> concretize matches the equivalent dense matrix expression, but Zygote sees an extra sensitivity contribution through the scaled-operator update path and returns a doubled gradient.

The first version of this PR avoided that by making ScaledOperator mutable, but review correctly pointed out that this is allocation-sensitive and not the right tradeoff. The revised fix keeps ScaledOperator immutable. Instead, after the scalar coefficient has been updated out-of-place, the updated ScalarOperator is converted to a private immutable _UpdatedScalarOperator snapshot before reconstructing the ScaledOperator.

That preserves the existing mutable ScalarOperator behavior for normal in-place coefficient updates while giving Zygote an immutable scalar value in the out-of-place path that is used by concretization. The in-place update_coefficients!(::ScaledOperator, ...) behavior is left unchanged, and this does not add new public API or dependencies.

The downgrade CI failure was a resolver conflict before tests ran: the root lower-bound environment selected Adapt v4.0.0, while the test dependency stack required compatibility with Adapt v4.5.2-4. Raising the package lower bound to ArrayInterface = "7.24" makes the downgraded root environment select an Adapt version compatible with the tests while keeping the downgrade check meaningful.

Tests

  • Focused Wrong result when Zygote differentiate through update_coefficient and concretize #305 reproducer comparing the operator and matrix gradients: operator and matrix gradients both evaluate to 1.0.
  • Focused Zygote and AD semantic test run:
    • julia --project=test -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate(); using Test; include("test/zygote.jl"); include("test/ad_semantics.jl")'
    • Zygote update_coefficients concretize scaled operator | 2 passed
    • AD semantic equivalence | 10 passed
  • Format check:
    • julia --project=@runic -e 'using Runic; exit(Runic.main(ARGS))' -- --check src/basic.jl src/scalar.jl test/runtests.jl test/ad_semantics.jl Project.toml
  • Full package tests:
    • julia --project=. -e 'using Pkg; Pkg.test()'
    • SciMLOperators | 896 passed, 2 broken, 898 total
  • Downgrade CI reproduction on Julia 1.10:
    • julia +1.10 --project=. -e 'import Pkg; Pkg.test(; coverage=true, julia_args=["--check-bounds=yes", "--compiled-modules=yes", "--depwarn=yes"], force_latest_compatible_version=false, allow_reresolve=false)'
    • SciMLOperators | 896 passed, 2 broken, 898 total

@AshtonSBradley AshtonSBradley marked this pull request as ready for review May 7, 2026 06:59
Comment thread src/basic.jl Outdated
(λ L)*(v) = λ * L(v)
"""
struct ScaledOperator{
mutable struct ScaledOperator{
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this will cause some allocations, seems like not the right solution.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I hope this is getting closer

@AshtonSBradley AshtonSBradley force-pushed the fix-zygote-concretize-gradient branch from dd10804 to 10e831a Compare May 8, 2026 20:36
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.

Wrong result when Zygote differentiate through update_coefficient and concretize

2 participants