implicit kwargs as an option, and verbose=false by default

Author: BatyLeo

ext/InferOptFrankWolfeExt.jl

@@ -8,6 +8,24 @@ using ImplicitDifferentiation: KrylovLinearSolver
 using InferOpt: InferOpt, RegularizedFrankWolfe
 using LinearAlgebra: dot
 
+"""
+    RegularizedFrankWolfe(linear_maximizer; Ω, Ω_grad, frank_wolfe_kwargs=(;), implicit_kwargs=(; linear_solver=KrylovLinearSolver(; verbose=false)))
+
+Construct a `RegularizedFrankWolfe` struct with a linear maximizer and the necessary components for the Frank-Wolfe algorithm.
+Set `implicit_kwargs` to `(; linear_solver=KrylovLinearSolver(; verbose=true))` if you want to see the solver potential warnings.
+"""
+function RegularizedFrankWolfe(
+    linear_maximizer;
+    Ω,
+    Ω_grad,
+    frank_wolfe_kwargs=NamedTuple(),
+    implicit_kwargs=(; linear_solver=KrylovLinearSolver(; verbose=false)),
+)
+    return RegularizedFrankWolfe(
+        linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs, implicit_kwargs
+    )
+end
+
 """
     LinearMaximizationOracleWithKwargs{F,K}
 Wraps a linear maximizer as a `FrankWolfe.LinearMinimizationOracle` with a sign switch and predefined keyword arguments.
@@ -42,12 +60,11 @@ function InferOpt.compute_probability_distribution(
     regularized::RegularizedFrankWolfe, θ::AbstractArray; kwargs...
 )
     shape = size(θ)
-    (; linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs) = regularized
+    (; linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs, implicit_kwargs) = regularized
     f(y, θ) = Ω(y) - dot(θ, y)
     f_grad1(y, θ) = Ω_grad(y) - θ
     maximizer(θ; shape, kwargs...) = vec(linear_maximizer(reshape(θ, shape); kwargs...))
     lmo = LinearMaximizationOracleWithKwargs(maximizer, (; shape, kwargs...))
-    implicit_kwargs = (; linear_solver=KrylovLinearSolver())
     dfw = DiffFW(f, f_grad1, lmo; implicit_kwargs)
     weights, atoms = dfw.implicit(vec(θ); frank_wolfe_kwargs=frank_wolfe_kwargs)
     probadist = FixedAtomsProbabilityDistribution(

src/layers/regularized/regularized_frank_wolfe.jl

@@ -7,14 +7,18 @@ Regularized optimization layer which relies on the Frank-Wolfe algorithm to defi
 ```
 
 !!! warning "Warning"
-    Since this is a conditional dependency, you need to have loaded the package DifferentiableFrankWolfe.jl before using `RegularizedFrankWolfe`.
+    Since this is a conditional dependency, you need to have loaded the following packages before using `RegularizedFrankWolfe`:
+    - `DifferentiableFrankWolfe.jl`
+    - `FrankWolfe.jl`
+    - `ImplicitDifferentiation.jl`
 
 # Fields
 
 - `linear_maximizer`: linear maximization oracle `θ -> argmax_{x ∈ C} θᵀx`, implicitly defines the polytope `C`
 - `Ω`: regularization function `Ω(y)`
 - `Ω_grad`: gradient function of the regularization function `∇Ω(y)`
 - `frank_wolfe_kwargs`: named tuple of keyword arguments passed to the Frank-Wolfe algorithm
+- `implicit_kwargs`: named tuple of keyword arguments passed to the implicit differentiation algorithm (in particular, the needed linear solver)
 
 # Frank-Wolfe parameters
 
@@ -29,24 +33,19 @@ Some values you can tune:
 
 See the documentation of FrankWolfe.jl for details.
 """
-struct RegularizedFrankWolfe{M,RF,RG,FWK} <: AbstractRegularized
+struct RegularizedFrankWolfe{M,RF,RG,FWK,IK} <: AbstractRegularized
     linear_maximizer::M
     Ω::RF
     Ω_grad::RG
     frank_wolfe_kwargs::FWK
-end
-
-"""
-    RegularizedFrankWolfe(linear_maximizer; Ω, Ω_grad, frank_wolfe_kwargs=(;))
-"""
-function RegularizedFrankWolfe(linear_maximizer; Ω, Ω_grad, frank_wolfe_kwargs=NamedTuple())
-    return RegularizedFrankWolfe(linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs)
+    implicit_kwargs::IK
 end
 
 function Base.show(io::IO, regularized::RegularizedFrankWolfe)
-    (; linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs) = regularized
+    (; linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs, implicit_kwargs) = regularized
     return print(
-        io, "RegularizedFrankWolfe($linear_maximizer, $Ω, $Ω_grad, $frank_wolfe_kwargs)"
+        io,
+        "RegularizedFrankWolfe($linear_maximizer, $Ω, $Ω_grad, $frank_wolfe_kwargs, $implicit_kwargs)",
     )
 end
 

test/InferOptTestUtils/src/pipeline.jl

@@ -21,7 +21,6 @@ function test_pipeline!(
     perf_storage = init_perf()
 
     for ep in 1:epochs
-        @info ep
         update_perf!(
             pl,
             perf_storage;