AugLag.jl

AugLag(; inner, kwargs...)

Augmented Lagrangian optimizer that solves constrained Optimization.jl problems by repeatedly calling an inner optimizer on the augmented objective.

Keyword Arguments

• inner: Inner optimizer used for the augmented subproblems.
• τ: Constraint violation contraction threshold.
• γ: Penalty growth factor.
• ϵ, ϵ_primal, ϵ_dual: Convergence tolerances.
• ρ_init: Initial penalty parameter. If nothing, a problem-scaled value is chosen.
• inner_kwargs: Keyword arguments forwarded to the inner optimizer.

classify_constraints(lcons, ucons)

Partition constraint indices given the SciMLBase contract lcons[i] ≤ c_i(θ) ≤ ucons[i]:

• eq_inds: rows where lcons[i] == ucons[i] (equality).
• ineq_upper_inds: rows with lcons[i] != ucons[i] and finite ucons[i], contributing the penalty for c_i ≤ ucons[i].
• ineq_lower_inds: rows with lcons[i] != ucons[i] and finite lcons[i], contributing the penalty for lcons[i] ≤ c_i.

A two-sided inequality (both bounds finite, lcons[i] < ucons[i]) appears in bothineq_upper_inds and ineq_lower_inds and gets a separate Lagrange multiplier per side.

generate_auglag(cache, eq_inds, ineq_upper_inds, ineq_lower_inds,
                λ, μ_upper, μ_lower, ρ_ref)

Build the augmented-Lagrangian subproblem function as an OptimizationFunction with analytical value, gradient, and fg! derived from the user's loss and constraints in cache.f.

The augmented Lagrangian is

L(θ; λ, μ_u, μ_l, ρ) = f(θ)
    + Σᵢ∈eq    [ λᵢ (cᵢ - lᵢ) + (ρ/2)(cᵢ - lᵢ)² ]
    + (1/(2ρ)) Σᵢ∈up max(0, μ_uᵢ + ρ (cᵢ - uᵢ))²
    + (1/(2ρ)) Σᵢ∈lo max(0, μ_lᵢ + ρ (lᵢ - cᵢ))²

so its gradient (used closed-form, not by AD'ing through L) is

∇L = ∇f
    + Σᵢ∈eq    (λᵢ + ρ (cᵢ - lᵢ)) ∇cᵢ
    + Σᵢ∈up    max(0, μ_uᵢ + ρ (cᵢ - uᵢ)) ∇cᵢ
    - Σᵢ∈lo    max(0, μ_lᵢ + ρ (lᵢ - cᵢ)) ∇cᵢ.

λ, μ_upper, μ_lower are mutable vectors and ρ_ref is a Ref{<:Real} (or any zero-arg-getindex container). The closures dereference them at call time, so the outer AugLag loop can update multipliers and the penalty in place between inner solves without rebuilding the function.

The constraint term of ∇L is a single vector-Jacobian product Jᵀv with cache.f.cons_vjp, where v holds the multiplier weight of every constraint row (zero for inactive inequalities). AugLag declares requiresconsvjp, so OptimizationBase always provides one: the user's, one built from the user's cons_j, the AD backend's pullback, or a chunked/colored AD Jacobian for forward-mode backends. Because inactive rows enter with a zero weight rather than being skipped, a non-finite derivative in an inactive row propagates to the gradient.

cons_tmp, v and Jᵀv are preallocated once with element type eltype(cache.u0). This is safe because the analytical gradient does not AD through this function — cache.f.grad and cache.f.cons_vjp handle their own AD internally.

Constraints and the data-iterator p

Constraints are treated as non-stochastic / batch-independent: the user's cons!(res, θ, p) is always invoked with the fullcache.p, regardless of which inner-solve batch is currently being processed. For the typical non-DataLoader case cache.p is just the user's p. For a DataLoaderp, cons! receives the iterator itself, and the user's constraint body may pull the underlying full data from it (e.g. via p.data for an MLUtils.DataLoader) — but the constraint is still a deterministic function of θ and the full data, never of a single batch.

Consistent with this, cache.f.cons_vjp is invoked without p and uses the p that was closed over at instantiation (the first batch for a data iterator). Since the constraint is by contract batch-independent, that closed-over p is irrelevant to its value.