Skip to content

Commit dbd9508

Browse files
Add test for Sophia + ComponentArrays + Enzyme shadow generation fix
This test verifies that the Sophia optimizer works correctly with ComponentArrays and Enzyme autodiff without encountering the shadow generation MethodError that was previously occurring. The test uses the existing minibatch test setup with ComponentArrays and creates an AutoEnzyme optimization function to ensure the fix works in the specific combination that was failing. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent edb45ad commit dbd9508

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

test/minibatch.jl

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,15 @@ res1 = Optimization.solve(optprob, Optimisers.Adam(0.05),
7171
callback = callback, maxiters = numEpochs)
7272
@test 10res1.objective < l1
7373

74+
# Test Sophia with ComponentArrays + Enzyme (shadow generation fix)
75+
optfun_enzyme = OptimizationFunction(loss_adjoint, Optimization.AutoEnzyme())
76+
optprob_enzyme = OptimizationProblem(optfun_enzyme, pp, train_loader)
77+
78+
res_sophia_enzyme = Optimization.solve(optprob_enzyme,
79+
Optimization.Sophia=0.01, k=5), callback = callback,
80+
maxiters = 50)
81+
@test res_sophia_enzyme.objective < Inf # Test that it runs without MethodError
82+
7483
optfun = OptimizationFunction(
7584
(θ, p) -> loss_adjoint(θ, batch,
7685
time_batch),

0 commit comments

Comments
 (0)