Skip to content

Commit a5fbd3d

Browse files
Move Sophia + ComponentArrays + Enzyme test to proper location
Moved the shadow generation fix test from minibatch.jl to native.jl where it belongs. The test now uses a simple Rosenbrock function with ComponentArrays rather than the complex neural ODE setup, making it more focused on testing the specific shadow generation compatibility issue that was fixed. Test location: test/native.jl (lines 65-74) - Simple ComponentVector with Rosenbrock function - Uses AutoEnzyme() to trigger the shadow generation path - Tests both optimization progress and successful completion 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent dbd9508 commit a5fbd3d

File tree

2 files changed

+11
-9
lines changed

2 files changed

+11
-9
lines changed

test/minibatch.jl

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -71,15 +71,6 @@ 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-
8374
optfun = OptimizationFunction(
8475
(θ, p) -> loss_adjoint(θ, batch,
8576
time_batch),

test/native.jl

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,3 +61,14 @@ optf1 = OptimizationFunction(loss, AutoSparseForwardDiff())
6161
prob1 = OptimizationProblem(optf1, rand(5), data)
6262
sol1 = solve(prob1, OptimizationOptimisers.Adam(), maxiters = 1000, callback = callback)
6363
@test sol1.objective < l0
64+
65+
# Test Sophia with ComponentArrays + Enzyme (shadow generation fix)
66+
using ComponentArrays
67+
x0_comp = ComponentVector(a = 0.0, b = 0.0)
68+
rosenbrock_comp(x, p = nothing) = (1 - x.a)^2 + 100 * (x.b - x.a^2)^2
69+
70+
optf_sophia = OptimizationFunction(rosenbrock_comp, AutoEnzyme())
71+
prob_sophia = OptimizationProblem(optf_sophia, x0_comp)
72+
res_sophia = solve(prob_sophia, Optimization.Sophia=0.01, k=5), maxiters = 50)
73+
@test res_sophia.objective < rosenbrock_comp(x0_comp) # Test optimization progress
74+
@test res_sophia.retcode == Optimization.SciMLBase.ReturnCode.Success

0 commit comments

Comments
 (0)