rework interface

Author: BatyLeo

src/imitation/fenchel_young_loss.jl

@@ -45,7 +45,7 @@ end
 
 function fenchel_young_loss_and_grad(
     fyl::FenchelYoungLoss{O}, θ::AbstractArray, y_true::AbstractArray; kwargs...
-) where {O<:AbstractRegularized{<:GeneralizedMaximizer}}
+) where {O<:AbstractRegularizedGeneralizedMaximizer}
     (; optimization_layer) = fyl
     ŷ = optimization_layer(θ; kwargs...)
     Ωy_true = compute_regularization(optimization_layer, y_true)

src/regularized/abstract_regularized.jl

@@ -1,7 +1,7 @@
 """
-    AbstractRegularized{parallel} <: AbstractOptimizationLayer
+    AbstractRegularized <: AbstractOptimizationLayer
 
-Convex regularization perturbation of a black box optimizer
+Convex regularization perturbation of a black box linear optimizer
 ```
 ŷ(θ) = argmax_{y ∈ C} {θᵀy - Ω(y)}
 ```
@@ -17,7 +17,29 @@ Convex regularization perturbation of a black box optimizer
 - [`SparseArgmax`](@ref)
 - [`RegularizedFrankWolfe`](@ref)
 """
-abstract type AbstractRegularized{O} <: AbstractOptimizationLayer end
+abstract type AbstractRegularized <: AbstractOptimizationLayer end
+
+"""
+    AbstractRegularizedGeneralizedMaximizer <: AbstractRegularized
+
+Convex regularization perturbation of a black box **generalized** optimizer
+```
+ŷ(θ) = argmax_{y ∈ C} {θᵀg(y) + h(y) - Ω(y)}
+with g and h functions of y.
+```
+
+# Interface
+
+- `(regularized::AbstractRegularized)(θ; kwargs...)`: return `ŷ(θ)`
+- `compute_regularization(regularized, y)`: return `Ω(y)`
+
+# Available implementations
+
+- [`SoftArgmax`](@ref)
+- [`SparseArgmax`](@ref)
+- [`RegularizedFrankWolfe`](@ref)
+"""
+abstract type AbstractRegularizedGeneralizedMaximizer <: AbstractRegularized end
 
 """
     compute_regularization(regularized, y)
@@ -26,9 +48,16 @@ Return the convex penalty `Ω(y)` associated with an `AbstractRegularized` layer
 """
 function compute_regularization end
 
+@required AbstractRegularized begin
+    #(regularized::AbstractRegularized)(θ::AbstractArray; kwargs...)
+    compute_regularization(::AbstractRegularized, ::AbstractArray)
+end
+
+"""
+TODO
+"""
 function get_maximizer end
 
-@required AbstractRegularized begin
-    compute_regularization(::AbstractRegularized, ::Any)
-    get_maximizer(::AbstractRegularized)
+@required AbstractRegularizedGeneralizedMaximizer begin
+    get_maximizer(::AbstractRegularizedGeneralizedMaximizer)
 end

src/regularized/regularized_frank_wolfe.jl

@@ -29,7 +29,7 @@ Some values you can tune:
 
 See the documentation of FrankWolfe.jl for details.
 """
-struct RegularizedFrankWolfe{M,RF,RG,FWK} <: AbstractRegularized{M}
+struct RegularizedFrankWolfe{M,RF,RG,FWK} <: AbstractRegularized
     linear_maximizer::M
     Ω::RF
     Ω_grad::RG

src/regularized/soft_argmax.jl

@@ -5,7 +5,7 @@ Soft argmax activation function `s(z) = (e^zᵢ / ∑ e^zⱼ)ᵢ`.
 
 Corresponds to regularized prediction on the probability simplex with entropic penalty.
 """
-struct SoftArgmax <: AbstractRegularized{nothing} end
+struct SoftArgmax <: AbstractRegularized end
 
 (::SoftArgmax)(z; kwargs...) = soft_argmax(z)
 compute_regularization(::SoftArgmax, y) = soft_argmax_regularization(y)

src/regularized/sparse_argmax.jl

@@ -5,7 +5,7 @@ Compute the Euclidean projection of the vector `z` onto the probability simplex.
 
 Corresponds to regularized prediction on the probability simplex with square norm penalty.
 """
-struct SparseArgmax <: AbstractRegularized{nothing} end
+struct SparseArgmax <: AbstractRegularized end
 
 (::SparseArgmax)(z; kwargs...) = sparse_argmax(z)
 compute_regularization(::SparseArgmax, y) = sparse_argmax_regularization(y)