Fix tests

src/InferOpt.jl

@@ -89,7 +89,7 @@ export SoftRank, soft_rank, soft_rank_l2, soft_rank_kl
 export SoftSort, soft_sort, soft_sort_l2, soft_sort_kl
 export RegularizedFrankWolfe
 
-export Perturbed
+export PerturbedOracle
 export PerturbedAdditive
 export PerturbedMultiplicative
 export LinearPerturbed

src/layers/perturbed/perturbed.jl

@@ -1,5 +1,5 @@
 """
-    Perturbed{D,F} <: AbstractOptimizationLayer
+    PerturbedOracle{D,F} <: AbstractOptimizationLayer
 
 Differentiable perturbation of a black box optimizer of type `F`, with perturbation of type `D`.
 
@@ -10,23 +10,25 @@ There are three different available constructors that behave differently in the
 - [`PerturbedAdditive`](@ref)
 - [`PerturbedMultiplicative`](@ref)
 """
-struct Perturbed{D,F,t,variance_reduction,G,R,S} <: AbstractOptimizationLayer
+struct PerturbedOracle{D,F,t,variance_reduction,G,R,S} <: AbstractOptimizationLayer
     reinforce::Reinforce{t,variance_reduction,F,D,G,R,S}
 end
 
-function (perturbed::Perturbed)(θ::AbstractArray; kwargs...)
+function (perturbed::PerturbedOracle)(θ::AbstractArray; kwargs...)
     return perturbed.reinforce(θ; kwargs...)
 end
 
-function get_maximizer(perturbed::Perturbed)
+function get_maximizer(perturbed::PerturbedOracle)
     return perturbed.reinforce.f
 end
 
-function compute_probability_distribution(perturbed::Perturbed, θ::AbstractArray; kwargs...)
+function compute_probability_distribution(
+    perturbed::PerturbedOracle, θ::AbstractArray; kwargs...
+)
     return empirical_distribution(perturbed.reinforce, θ; kwargs...)
 end
 
-function Base.show(io::IO, perturbed::Perturbed{<:AbstractPerturbation})
+function Base.show(io::IO, perturbed::PerturbedOracle{<:AbstractPerturbation})
     (; reinforce) = perturbed
     nb_samples = reinforce.nb_samples
     ε = reinforce.dist_constructor.ε
@@ -36,24 +38,36 @@ function Base.show(io::IO, perturbed::Perturbed{<:AbstractPerturbation})
     f = reinforce.f
     return print(
         io,
-        "Perturbed($f, ε=$ε, nb_samples=$nb_samples, perturbation=$perturbation, rng=$(typeof(rng)), seed=$seed)",
+        "PerturbedOracle($f, ε=$ε, nb_samples=$nb_samples, perturbation=$perturbation, rng=$(typeof(rng)), seed=$seed)",
     )
 end
 
 """
 doc
 """
-function LinearPerturbed(
+function PerturbedOracle(
     maximizer,
-    dist_constructor,
-    dist_logdensity_grad=nothing;
-    g=nothing,
-    h=nothing,
+    dist_constructor;
+    dist_logdensity_grad=nothing,
+    nb_samples=1,
+    variance_reduction=true,
+    threaded=false,
+    seed=nothing,
+    rng=Random.default_rng(),
     kwargs...,
 )
-    linear_maximizer = LinearMaximizer(; maximizer, g, h)
-    return Perturbed(
-        Reinforce(linear_maximizer, dist_constructor, dist_logdensity_grad; kwargs...)
+    return PerturbedOracle(
+        Reinforce(
+            maximizer,
+            dist_constructor,
+            dist_logdensity_grad;
+            nb_samples,
+            variance_reduction,
+            threaded,
+            seed,
+            rng,
+            kwargs...,
+        ),
     )
 end
 
@@ -69,26 +83,22 @@ function PerturbedAdditive(
     seed=nothing,
     threaded=false,
     rng=Random.default_rng(),
-    g=identity_kw,
-    h=zero ∘ eltype_kw,
     dist_logdensity_grad=if (perturbation_dist == Normal(0, 1))
         (η, θ) -> ((η .- θ) ./ ε^2,)
     else
         nothing
     end,
 )
     dist_constructor = AdditivePerturbation(perturbation_dist, float(ε))
-    return LinearPerturbed(
+    return PerturbedOracle(
         maximizer,
-        dist_constructor,
-        dist_logdensity_grad;
+        dist_constructor;
+        dist_logdensity_grad,
         nb_samples,
         variance_reduction,
         seed,
         threaded,
         rng,
-        g,
-        h,
     )
 end
 
@@ -104,25 +114,21 @@ function PerturbedMultiplicative(
     seed=nothing,
     threaded=false,
     rng=Random.default_rng(),
-    g=identity_kw,
-    h=zero ∘ eltype_kw,
     dist_logdensity_grad=if (perturbation_dist == Normal(0, 1))
         (η, θ) -> (inv.(ε^2 .* θ) .* (η .- θ),)
     else
         nothing
     end,
 )
     dist_constructor = MultiplicativePerturbation(perturbation_dist, float(ε))
-    return LinearPerturbed(
+    return PerturbedOracle(
         maximizer,
-        dist_constructor,
-        dist_logdensity_grad;
+        dist_constructor;
+        dist_logdensity_grad,
         nb_samples,
         variance_reduction,
         seed,
         threaded,
         rng,
-        g,
-        h,
     )
 end

src/losses/fenchel_young_loss.jl

@@ -59,7 +59,10 @@ function fenchel_young_loss_and_grad(
 end
 
 function fenchel_young_loss_and_grad(
-    fyl::FenchelYoungLoss{<:Perturbed}, θ::AbstractArray, y_true::AbstractArray; kwargs...
+    fyl::FenchelYoungLoss{<:PerturbedOracle},
+    θ::AbstractArray,
+    y_true::AbstractArray;
+    kwargs...,
 )
     (; optimization_layer) = fyl
     maximizer = get_maximizer(optimization_layer)
@@ -82,7 +85,7 @@ end
 ## Specific overrides for perturbed layers
 
 function fenchel_young_F_and_first_part_of_grad(
-    perturbed::Perturbed{<:AdditivePerturbation}, θ::AbstractArray; kwargs...
+    perturbed::PerturbedOracle{<:AdditivePerturbation}, θ::AbstractArray; kwargs...
 )
     (; reinforce) = perturbed
     maximizer = get_maximizer(perturbed)
@@ -98,7 +101,7 @@ function fenchel_young_F_and_first_part_of_grad(
 end
 
 function fenchel_young_F_and_first_part_of_grad(
-    perturbed::Perturbed{<:MultiplicativePerturbation}, θ::AbstractArray; kwargs...
+    perturbed::PerturbedOracle{<:MultiplicativePerturbation}, θ::AbstractArray; kwargs...
 )
     (; reinforce) = perturbed
     maximizer = get_maximizer(perturbed)

src/utils/linear_maximizer.jl

@@ -18,6 +18,10 @@ function Base.show(io::IO, f::LinearMaximizer)
     return print(io, "LinearMaximizer($maximizer, $g, $h)")
 end
 
+function LinearMaximizer(maximizer; g=identity_kw, h=zero ∘ eltype_kw)
+    return LinearMaximizer(maximizer, g, h)
+end
+
 # Callable calls the wrapped maximizer
 function (f::LinearMaximizer)(θ::AbstractArray; kwargs...)
     return f.maximizer(θ; kwargs...)

test/learning_generalized_maximizer.jl

@@ -13,7 +13,7 @@
 
     @test y == [1 0 1; 0 1 0; 1 1 1]
 
-    generalized_maximizer = LinearMaximizer(; maximizer=max_pricing, g, h)
+    generalized_maximizer = LinearMaximizer(max_pricing; g, h)
 
     @test generalized_maximizer(θ; instance) == y
 
@@ -29,8 +29,8 @@ end
 
     true_encoder = encoder_factory()
 
-    perturbed = PerturbedAdditive(max_pricing; ε=1.0, nb_samples=10, g, h)
-    maximizer = InferOpt.get_maximizer(perturbed)
+    maximizer = LinearMaximizer(max_pricing; g, h)
+    perturbed = PerturbedAdditive(maximizer; ε=1.0, nb_samples=10)
     function cost(y; instance)
         return -objective_value(maximizer, true_encoder(instance), y; instance)
     end
@@ -55,8 +55,8 @@ end
 
     true_encoder = encoder_factory()
 
-    perturbed = PerturbedMultiplicative(max_pricing; ε=1.0, nb_samples=10, g, h)
-    maximizer = InferOpt.get_maximizer(perturbed)
+    maximizer = LinearMaximizer(max_pricing; g, h)
+    perturbed = PerturbedMultiplicative(maximizer; ε=1.0, nb_samples=10)
     function cost(y; instance)
         return -objective_value(maximizer, true_encoder(instance), y; instance)
     end
@@ -80,8 +80,10 @@ end
 
     true_encoder = encoder_factory()
 
-    perturbed = PerturbedAdditive(max_pricing; ε=1.0, nb_samples=10, g, h)
-    maximizer = InferOpt.get_maximizer(perturbed)
+    maximizer = LinearMaximizer(max_pricing; g, h)
+    @info maximizer g h
+    perturbed = PerturbedAdditive(maximizer; ε=1.0, nb_samples=10)
+    @info perturbed
     function cost(y; instance)
         return -objective_value(maximizer, true_encoder(instance), y; instance)
     end
@@ -106,8 +108,8 @@ end
 
     true_encoder = encoder_factory()
 
-    perturbed = PerturbedMultiplicative(max_pricing; ε=0.1, nb_samples=10, g, h)
-    maximizer = InferOpt.get_maximizer(perturbed)
+    maximizer = LinearMaximizer(max_pricing; g, h)
+    perturbed = PerturbedMultiplicative(maximizer; ε=0.1, nb_samples=10)
     function cost(y; instance)
         return -objective_value(maximizer, true_encoder(instance), y; instance)
     end
@@ -180,7 +182,7 @@ end
 
     true_encoder = encoder_factory()
 
-    generalized_maximizer = GeneralizedMaximizer(max_pricing; g, h)
+    generalized_maximizer = LinearMaximizer(max_pricing; g, h)
     function cost(y; instance)
         return -objective_value(generalized_maximizer, true_encoder(instance), y; instance)
     end
@@ -207,7 +209,7 @@ end
 
     true_encoder = encoder_factory()
 
-    generalized_maximizer = GeneralizedMaximizer(max_pricing; g, h)
+    generalized_maximizer = LinearMaximizer(max_pricing; g, h)
     function cost(y; instance)
         return -objective_value(generalized_maximizer, true_encoder(instance), y; instance)
     end

test/paths.jl

@@ -155,7 +155,10 @@ end
         maximizer=identity_kw,
         loss=FenchelYoungLoss(
             PerturbedAdditive(
-                shortest_path_maximizer; ε=1.0, nb_samples=5, perturbation=LogNormal(0, 1)
+                shortest_path_maximizer;
+                ε=1.0,
+                nb_samples=5,
+                perturbation_dist=LogNormal(0, 1),
             ),
         ),
         error_function=mse_kw,

test/perturbed.jl

@@ -50,7 +50,7 @@ end
     p(θ) = MvNormal(θ, ε^2 * I)
     oracle(η) = η
 
-    po = PerturbedOracle(oracle, p; nb_samples=1_000, seed=0) # TODO: fix this
+    po = PerturbedOracle(oracle, p; nb_samples=1_000, seed=0)
     pa = PerturbedAdditive(oracle; ε, nb_samples=1_000, seed=0)
 
     θ = randn(10)