Skip to content

Commit 051a39e

Browse files
Fix Sophia optimizer shadow generation for ComponentArrays
This fix resolves the issue where Sophia optimizer fails when used with ComponentArrays and DifferentiationInterface/Enzyme autodiff due to incorrect shadow type generation. Problem: - Line 154 in sophia.jl used `randn(uType, length(θ))` which creates a standard Vector, losing the structure of θ when it's a ComponentVector - This causes MethodError when Enzyme tries to create Duplicated types with mismatched structures Solution: - Replace with `u = similar(θ); randn!(u)` pattern - This preserves the structure of θ while filling with random values - Follows the same pattern used elsewhere in the codebase with `zero(θ)` References: https://discourse.julialang.org/t/differentiationinterface-not-generating-correct-shadow/131868 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent 8f8e4ff commit 051a39e

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/sophia.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,8 @@ function SciMLBase.__solve(cache::OptimizationCache{
151151

152152
if i % cache.opt.k == 1
153153
hₜ₋₁ = copy(hₜ)
154-
u = randn(uType, length(θ))
154+
u = similar(θ)
155+
randn!(u)
155156
f.hv(hₜ, θ, u, d)
156157
hₜ = βs[2] .* hₜ₋₁ + (1 - βs[2]) .* (u .* hₜ)
157158
end

0 commit comments

Comments
 (0)