Always pass p explicitly to instantiated functions (stale p after reinit!, data iterators silently use the first batch) #1383

Summary

instantiate_function bakes p into the derivative and constraint closures (grad, fg, hess, hv, cons, cons_j, lag_h, ...) when the cache is built. The objective f.f is left unwrapped, so solvers must pass p to it. Solvers call the derivatives both ways, grad(G, θ) (baked p) and grad(G, θ, p) (explicit p), and the backends don't agree on which forms they support. This causes three problems:

1. reinit!(cache; p = ...) is silently ignored by every baked closure. reinit! only sets cache.reinit_cache.p and does not re-instantiate anything.
2. With a data iterator as p, the closures are baked with the first batch (OptimizationBase/src/cache.jl:84-90), while cache.p holds the whole iterator. Solvers that don't handle iterators silently fit batch 1.
3. Backends are inconsistent about whether a passed p is used, ignored, or appended as an extra argument (details below).

Proposal: always pass p explicitly. Every instantiated closure takes the parameters as an argument, solvers always pass cache.p (or the current batch), and nothing depends on the p seen at instantiation. Data iterators need a decision first; see the last section.

Reproductions

All runs below use master (bf6d75a) plus #1381, on Julia 1.12.7.

Stale p after reinit!

using OptimizationBase, OptimizationOptimisers, Optimisers, ForwardDiff
using SciMLBase: OptimizationFunction, OptimizationProblem
objective(x, p) = (p[1] - x[1])^2
grad!(G, x, p) = (G[1] = -2 * (p[1] - x[1]))
for optf in (OptimizationFunction(objective, OptimizationBase.AutoForwardDiff()),
             OptimizationFunction(objective; grad = grad!))
    cache = OptimizationBase.init(OptimizationProblem(optf, zeros(1), [1.0]), Optimisers.Adam(0.1), maxiters = 1000)
    OptimizationBase.solve!(cache)
    cache = OptimizationBase.reinit!(cache; p = [2.0])
    @show OptimizationBase.solve!(cache).u cache.p
end

This is the existing @test_broken in lib/OptimizationOptimisers/test/core_tests.jl:99-101. Related @test_brokens: lib/OptimizationMOI/test/core_tests.jl:251,263 ("needs reinit/remake fixes") and lib/OptimizationIpopt/test/core_tests.jl:192-202 (MTK). Ipopt's other reinit test passes because Ipopt forwards p explicitly when it can (_accepts_live_p).

A solver that doesn't handle iterators fits only batch 1

using OptimizationBase, OptimizationOptimJL, ForwardDiff, MLUtils
using SciMLBase: OptimizationFunction, OptimizationProblem
data = MLUtils.DataLoader(([1.0, 1.0, 5.0, 5.0],), batchsize = 2)
loss(u, batch) = sum(abs2, u[1] .- batch[1])
optf = OptimizationFunction(loss, OptimizationBase.AutoForwardDiff())
solve(OptimizationProblem(optf, [0.0], data), LBFGS(), maxiters = 50).u  # [1.0]; the full-data mean is 3.0

There is no error or warning. The objective only ever receives the first batch, because all evaluations go through the closures baked with batch 1.

Data iterators: current handling

• cache.jl:84-90 instantiates the closures from ReInitCache(u0, iterate(prob.p)[1]) and stores the iterator in cache.p. reinit! never updates the first-batch copy.
• Iterator-aware solvers:
  • OptimizationOptimisers and OptimizationSophia pass the batch d explicitly in their iterator branches (objective, grad, fg, and hv in Sophia). That part works: with save_best = false, Adam on the DataLoader above converges to the full-data mean.
  • OptimizationAuglag passes the whole iterator to cons (by design, see auglag_function.jl:65-80). cons_j is baked with batch 1, which is documented in test/core_tests.jl:339. The inner problem gets the iterator as p, so if the inner solver isn't iterator-aware, the objective gets the whole iterator while the derivatives use batch 1 (from reading the code; not run).
• All other solvers ignore iterators. They pass cache.p (the whole iterator) to the objective, while baked derivatives use batch 1. Solvers that pass p explicitly to derivatives (Ipopt, SciPy, ODE, the NLPModels grad) hand the whole iterator to them.
• Side finding: with minibatches, save_best in OptimizationOptimisers compares single-batch losses from different batches. The Adam run above returns u = [1.03] with the default save_best = true, and [3.0] with save_best = false.

With "always pass p", the p slot holds the batch in iterator-aware solvers. Solvers that aren't iterator-aware should reject a data iterator (for example through a trait or a check in __solve) instead of silently using batch 1. We also need to decide what instantiate_function sees for an iterator: probably just a prototype batch for preparation, never a value that gets used.

Backend inconsistencies

Verified by running:

• NoAD's if p != NullParameters() guards around the 3-arg methods (function.jl) don't do anything. Julia hoists local method definitions, so grad(G, x, p) exists even for NullParameters(). Both arities always exist; the guards are dead code. The same pattern appears in OptimizationDISparseExt.jl.

From reading the code (not run):

The MTK objective is mtkprob.f.f, which expects mtkprob.p, while solvers pass cache.p.

Solver call sites (from reading the code)

The objective always gets p explicitly, after #1381, which fixed the only exceptions: the grad-only branches of Optimisers and Sophia. For derivatives and constraints:

• Baked p, so stale after reinit! and batch 1 with iterators:
  • OptimizationOptimisers (fg 122, grad 127 in the non-iterator branch)
  • OptimizationSophia (128, 130)
  • OptimizationOptimJL (all derivative and constraint sites)
  • OptimizationLBFGSB, OptimizationNLopt, OptimizationManopt, OptimizationGCMAES, OptimizationSpeedMapping
  • NLPModels / MadNLP (hess, lag_h, all constraints)
  • OptimizationMOI nlp.jl (all derivatives and constraints)
  • Ipopt cons_h (evaluator.jl:239)
  • constraints in OptimizationODE, OptimizationEvolutionary, OptimizationPRIMA
• Explicit p:
  • Ipopt (when _accepts_live_p), SciPy, ODE (derivatives), NLPModels grad, Auglag
  • SimpleOptimization builds its gradient from Fix2(cache.f.f, cache.p) at solve time
• Explicit p can't work with MTK closures (the extra-argument issue): NLPModels grad, SciPy. Ipopt excludes AutoSymbolics from its explicit-p path.

Related

• Pass parameters to the objective when only grad is supplied (Optimisers, Sophia) #1381: objective called without p in the grad-only branch.
• Clarify data iterator interface #1203: clarify the data iterator interface.

Written by Claude Code on Sebastian's behalf.

🤖 Generated with Claude Code · model: claude-opus-5-5[1m]
Review: unreviewed