Remove all <:Real type constraints (#71)

gdalle · web-flow · commit b9f84b973fd1 · 2023-06-28T11:58:19.000+02:00
diff --git a/ext/InferOptFrankWolfeExt.jl b/ext/InferOptFrankWolfeExt.jl
@@ -8,7 +8,7 @@ using LinearAlgebra: dot
 ## Forward pass
 
 function InferOpt.compute_probability_distribution(
-    dfw::DiffFW, θ::AbstractArray{<:Real}; frank_wolfe_kwargs=NamedTuple()
+    dfw::DiffFW, θ::AbstractArray; frank_wolfe_kwargs=NamedTuple()
 )
     weights, atoms = dfw.implicit(θ; frank_wolfe_kwargs=frank_wolfe_kwargs)
     probadist = FixedAtomsProbabilityDistribution(atoms, weights)
@@ -23,7 +23,7 @@ Construct a `DifferentiableFrankWolfe.DiffFW` struct and call `compute_probabili
 Keyword arguments are passed to the underlying linear maximizer.
 """
 function InferOpt.compute_probability_distribution(
-    regularized::RegularizedGeneric, θ::AbstractArray{<:Real}; kwargs...
+    regularized::RegularizedGeneric, θ::AbstractArray; kwargs...
 )
     (; maximizer, Ω, Ω_grad, frank_wolfe_kwargs) = regularized
     f(y, θ) = Ω(y) - dot(θ, y)
@@ -41,7 +41,7 @@ Apply `compute_probability_distribution(regularized, θ)` and return the expecta
 
 Keyword arguments are passed to the underlying linear maximizer.
 """
-function (regularized::RegularizedGeneric)(θ::AbstractArray{<:Real}; kwargs...)
+function (regularized::RegularizedGeneric)(θ::AbstractArray; kwargs...)
     probadist = compute_probability_distribution(regularized, θ; kwargs...)
     return compute_expectation(probadist)
 end
diff --git a/src/fenchel_young/fenchel_young.jl b/src/fenchel_young/fenchel_young.jl
@@ -19,18 +19,13 @@ end
 
 ## Forward pass
 
-function (fyl::FenchelYoungLoss)(
-    θ::AbstractArray{<:Real}, y_true::AbstractArray{<:Real}; kwargs...
-)
+function (fyl::FenchelYoungLoss)(θ::AbstractArray, y_true::AbstractArray; kwargs...)
     l, _ = fenchel_young_loss_and_grad(fyl, θ, y_true; kwargs...)
     return l
 end
 
 @traitfn function fenchel_young_loss_and_grad(
-    fyl::FenchelYoungLoss{P},
-    θ::AbstractArray{<:Real},
-    y_true::AbstractArray{<:Real};
-    kwargs...,
+    fyl::FenchelYoungLoss{P}, θ::AbstractArray, y_true::AbstractArray; kwargs...
 ) where {P; IsRegularized{P}}
     (; predictor) = fyl
     ŷ = predictor(θ; kwargs...)
@@ -42,10 +37,7 @@ end
 end
 
 function fenchel_young_loss_and_grad(
-    fyl::FenchelYoungLoss{P},
-    θ::AbstractArray{<:Real},
-    y_true::AbstractArray{<:Real};
-    kwargs...,
+    fyl::FenchelYoungLoss{P}, θ::AbstractArray, y_true::AbstractArray; kwargs...
 ) where {P<:AbstractPerturbed}
     (; predictor) = fyl
     F, almost_ŷ = fenchel_young_F_and_first_part_of_grad(predictor, θ; kwargs...)
@@ -57,10 +49,7 @@ end
 ## Backward pass
 
 function ChainRulesCore.rrule(
-    fyl::FenchelYoungLoss,
-    θ::AbstractArray{<:Real},
-    y_true::AbstractArray{<:Real};
-    kwargs...,
+    fyl::FenchelYoungLoss, θ::AbstractArray, y_true::AbstractArray; kwargs...
 )
     l, g = fenchel_young_loss_and_grad(fyl, θ, y_true; kwargs...)
     fyl_pullback(dl) = NoTangent(), dl * g, NoTangent()
diff --git a/src/fenchel_young/perturbed.jl b/src/fenchel_young/perturbed.jl
@@ -1,5 +1,5 @@
 function compute_F_and_y_samples(
-    perturbed::AbstractPerturbed{false}, θ::AbstractArray{<:Real}, Z_samples; kwargs...
+    perturbed::AbstractPerturbed{false}, θ::AbstractArray, Z_samples; kwargs...
 )
     F_and_y_samples = [
         fenchel_young_F_and_first_part_of_grad(perturbed, θ, Z; kwargs...) for
@@ -9,26 +9,23 @@ function compute_F_and_y_samples(
 end
 
 function compute_F_and_y_samples(
-    perturbed::AbstractPerturbed{true}, θ::AbstractArray{<:Real}, Z_samples; kwargs...
+    perturbed::AbstractPerturbed{true}, θ::AbstractArray, Z_samples; kwargs...
 )
     return ThreadsX.map(
         Z -> fenchel_young_F_and_first_part_of_grad(perturbed, θ, Z; kwargs...), Z_samples
     )
 end
 
 function fenchel_young_F_and_first_part_of_grad(
-    perturbed::AbstractPerturbed, θ::AbstractArray{<:Real}; kwargs...
+    perturbed::AbstractPerturbed, θ::AbstractArray; kwargs...
 )
     Z_samples = sample_perturbations(perturbed, θ)
     F_and_y_samples = compute_F_and_y_samples(perturbed, θ, Z_samples; kwargs...)
     return mean(first, F_and_y_samples), mean(last, F_and_y_samples)
 end
 
 function fenchel_young_F_and_first_part_of_grad(
-    perturbed::PerturbedAdditive,
-    θ::AbstractArray{<:Real},
-    Z::AbstractArray{<:Real};
-    kwargs...,
+    perturbed::PerturbedAdditive, θ::AbstractArray, Z::AbstractArray; kwargs...
 )
     (; maximizer, ε) = perturbed
     θ_perturbed = θ .+ ε .* Z
@@ -38,10 +35,7 @@ function fenchel_young_F_and_first_part_of_grad(
 end
 
 function fenchel_young_F_and_first_part_of_grad(
-    perturbed::PerturbedMultiplicative,
-    θ::AbstractArray{<:Real},
-    Z::AbstractArray{<:Real};
-    kwargs...,
+    perturbed::PerturbedMultiplicative, θ::AbstractArray, Z::AbstractArray; kwargs...
 )
     (; maximizer, ε) = perturbed
     eZ = exp.(ε .* Z .- ε^2)
diff --git a/src/imitation_loss/imitation_loss.jl b/src/imitation_loss/imitation_loss.jl
@@ -53,7 +53,7 @@ Retrieve `y_true` from `t_true`. `t_true` must contain an `y_true` field.
 """
 get_y_true(t_true::NamedTuple) = t_true.y_true
 
-function prediction_and_loss(l::ImitationLoss, θ::AbstractArray{<:Real}, t_true; kwargs...)
+function prediction_and_loss(l::ImitationLoss, θ::AbstractArray, t_true; kwargs...)
     (; base_loss, Ω, maximizer, α) = l
     y_true = get_y_true(t_true)
     ŷ = maximizer(θ, t_true; kwargs...)
@@ -63,14 +63,14 @@ end
 
 ## Forward pass
 
-function (l::ImitationLoss)(θ::AbstractArray{<:Real}, t_true; kwargs...)
+function (l::ImitationLoss)(θ::AbstractArray, t_true; kwargs...)
     _, l = prediction_and_loss(l, θ, t_true; kwargs...)
     return l
 end
 
 ## Backward pass
 
-function ChainRulesCore.rrule(l::ImitationLoss, θ::AbstractArray{<:Real}, t_true; kwargs...)
+function ChainRulesCore.rrule(l::ImitationLoss, θ::AbstractArray, t_true; kwargs...)
     (; α) = l
     y_true = get_y_true(t_true)
     ŷ, l = prediction_and_loss(l, θ, t_true; kwargs...)
diff --git a/src/interpolation/interpolation.jl b/src/interpolation/interpolation.jl
@@ -21,13 +21,11 @@ end
 
 Interpolation(maximizer; λ=1.0) = Interpolation(maximizer, float(λ))
 
-function (interpolation::Interpolation)(θ::AbstractArray{<:Real}; kwargs...)
+function (interpolation::Interpolation)(θ::AbstractArray; kwargs...)
     return interpolation.maximizer(θ; kwargs...)
 end
 
-function ChainRulesCore.rrule(
-    interpolation::Interpolation, θ::AbstractArray{<:Real}; kwargs...
-)
+function ChainRulesCore.rrule(interpolation::Interpolation, θ::AbstractArray; kwargs...)
     (; maximizer, λ) = interpolation
     y = maximizer(θ; kwargs...)
     function interpolation_pullback(dy)
diff --git a/src/perturbed/abstract_perturbed.jl b/src/perturbed/abstract_perturbed.jl
@@ -29,7 +29,7 @@ abstract type AbstractPerturbed{parallel} end
 
 Draw random perturbations `Z` which will be applied to the objective direction `θ`.
 """
-function sample_perturbations(perturbed::AbstractPerturbed, θ::AbstractArray{<:Real})
+function sample_perturbations(perturbed::AbstractPerturbed, θ::AbstractArray)
     (; rng, seed, nb_samples) = perturbed
     seed!(rng, seed)
     Z_samples = [randn(rng, size(θ)) for _ in 1:nb_samples]
@@ -38,26 +38,26 @@ end
 
 function compute_atoms(
     perturbed::AbstractPerturbed{false},
-    θ::AbstractArray{<:Real},
-    Z_samples::Vector{<:AbstractArray{<:Real}};
+    θ::AbstractArray,
+    Z_samples::Vector{<:AbstractArray};
     kwargs...,
 )
     return [perturb_and_optimize(perturbed, θ, Z; kwargs...) for Z in Z_samples]
 end
 
 function compute_atoms(
     perturbed::AbstractPerturbed{true},
-    θ::AbstractArray{<:Real},
-    Z_samples::Vector{<:AbstractArray{<:Real}};
+    θ::AbstractArray,
+    Z_samples::Vector{<:AbstractArray};
     kwargs...,
 )
     return ThreadsX.map(Z -> perturb_and_optimize(perturbed, θ, Z; kwargs...), Z_samples)
 end
 
 function compute_probability_distribution(
     perturbed::AbstractPerturbed,
-    θ::AbstractArray{<:Real},
-    Z_samples::Vector{<:AbstractArray{<:Real}};
+    θ::AbstractArray,
+    Z_samples::Vector{<:AbstractArray};
     kwargs...,
 )
     atoms = compute_atoms(perturbed, θ, Z_samples; kwargs...)
@@ -72,7 +72,7 @@ end
 Turn random perturbations of `θ` into a distribution on polytope vertices.
 """
 function compute_probability_distribution(
-    perturbed::AbstractPerturbed, θ::AbstractArray{<:Real}; kwargs...
+    perturbed::AbstractPerturbed, θ::AbstractArray; kwargs...
 )
     Z_samples = sample_perturbations(perturbed, θ)
     return compute_probability_distribution(perturbed, θ, Z_samples; kwargs...)
@@ -83,7 +83,7 @@ end
 
 Apply `compute_probability_distribution(perturbed, θ)` and return the expectation.
 """
-function (perturbed::AbstractPerturbed)(θ::AbstractArray{<:Real}; kwargs...)
+function (perturbed::AbstractPerturbed)(θ::AbstractArray; kwargs...)
     probadist = compute_probability_distribution(perturbed, θ; kwargs...)
     return compute_expectation(probadist)
 end
diff --git a/src/perturbed/additive.jl b/src/perturbed/additive.jl
@@ -49,10 +49,7 @@ end
 ## Forward pass
 
 function perturb_and_optimize(
-    perturbed::PerturbedAdditive,
-    θ::AbstractArray{<:Real},
-    Z::AbstractArray{<:Real};
-    kwargs...,
+    perturbed::PerturbedAdditive, θ::AbstractArray, Z::AbstractArray; kwargs...
 )
     (; maximizer, ε) = perturbed
     θ_perturbed = θ .+ ε .* Z
@@ -65,7 +62,7 @@ end
 function ChainRulesCore.rrule(
     ::typeof(compute_probability_distribution),
     perturbed::PerturbedAdditive,
-    θ::AbstractArray{<:Real};
+    θ::AbstractArray;
     kwargs...,
 )
     (; ε) = perturbed
diff --git a/src/perturbed/multiplicative.jl b/src/perturbed/multiplicative.jl
@@ -51,10 +51,7 @@ end
 ## Forward pass
 
 function perturb_and_optimize(
-    perturbed::PerturbedMultiplicative,
-    θ::AbstractArray{<:Real},
-    Z::AbstractArray{<:Real};
-    kwargs...,
+    perturbed::PerturbedMultiplicative, θ::AbstractArray, Z::AbstractArray; kwargs...
 )
     (; maximizer, ε) = perturbed
     θ_perturbed = θ .* exp.(ε .* Z .- ε^2)
@@ -67,7 +64,7 @@ end
 function ChainRulesCore.rrule(
     ::typeof(compute_probability_distribution),
     perturbed::PerturbedMultiplicative,
-    θ::AbstractArray{<:Real};
+    θ::AbstractArray;
     kwargs...,
 )
     (; ε) = perturbed
diff --git a/src/plus_identity/plus_identity.jl b/src/plus_identity/plus_identity.jl
@@ -18,11 +18,11 @@ function Base.show(io::IO, plusid::PlusIdentity)
     return print(io, "PlusIdentity($(plusid.maximizer)")
 end
 
-function (plusid::PlusIdentity)(θ::AbstractArray{<:Real}; kwargs...)
+function (plusid::PlusIdentity)(θ::AbstractArray; kwargs...)
     return plusid.maximizer(θ; kwargs...)
 end
 
-function ChainRulesCore.rrule(plusid::PlusIdentity, θ::AbstractArray{<:Real}; kwargs...)
+function ChainRulesCore.rrule(plusid::PlusIdentity, θ::AbstractArray; kwargs...)
     y = plusid.maximizer(θ; kwargs...)
     plusid_pullback(dy) = NoTangent(), dy
     return y, plusid_pullback
diff --git a/src/regularized/regularized_generic.jl b/src/regularized/regularized_generic.jl
@@ -46,7 +46,7 @@ end
 
 @traitimpl IsRegularized{RegularizedGeneric}
 
-function compute_regularization(regularized::RegularizedGeneric, y::AbstractArray{<:Real})
+function compute_regularization(regularized::RegularizedGeneric, y::AbstractArray)
     return regularized.Ω(y)
 end
 
@@ -55,7 +55,7 @@ end
 
 Apply `compute_probability_distribution(regularized, θ, kwargs...)` and return the expectation.
 """
-function (regularized::RegularizedGeneric)(θ::AbstractArray{<:Real}; kwargs...)
+function (regularized::RegularizedGeneric)(θ::AbstractArray; kwargs...)
     probadist = compute_probability_distribution(regularized, θ; kwargs...)
     return compute_expectation(probadist)
 end
diff --git a/src/regularized/regularized_utils.jl b/src/regularized/regularized_utils.jl
@@ -24,7 +24,7 @@ isprobadist(p::AbstractVector{R}) where {R<:Real} = all(isproba, p) && sum(p) 
 
 Compute the squared Euclidean norm of `x` and divide it by 2.
 """
-function half_square_norm(x::AbstractArray{<:Real})
+function half_square_norm(x::AbstractArray)
     return sum(abs2, x) / 2
 end
 
@@ -43,7 +43,7 @@ function shannon_entropy(p::AbstractVector{R}) where {R<:Real}
     return H
 end
 
-negative_shannon_entropy(p::AbstractVector{<:Real}) = -shannon_entropy(p)
+negative_shannon_entropy(p::AbstractVector) = -shannon_entropy(p)
 
 """
     one_hot_argmax(z)
@@ -61,7 +61,7 @@ end
 
 Compute the vector `r` such that `rᵢ` is the rank of `θᵢ` in `θ`.
 """
-function ranking(θ::AbstractVector{<:Real}; rev::Bool=false, kwargs...)
+function ranking(θ::AbstractVector; rev::Bool=false, kwargs...)
     return invperm(sortperm(θ; rev=rev))
 end
 
diff --git a/src/regularized/soft_argmax.jl b/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.
 """
-function soft_argmax(z::AbstractVector{<:Real}; kwargs...)
+function soft_argmax(z::AbstractVector; kwargs...)
     s = exp.(z) / sum(exp, z)
     return s
 end
diff --git a/src/regularized/sparse_argmax.jl b/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.
 """
-function sparse_argmax(z::AbstractVector{<:Real}; kwargs...)
+function sparse_argmax(z::AbstractVector; kwargs...)
     p, _ = simplex_projection_and_support(z)
     return p
 end
@@ -25,7 +25,7 @@ Compute the Euclidean projection `p` of `z` on the probability simplex (also cal
 
 Reference: <https://arxiv.org/abs/1602.02068>.
 """
-function simplex_projection_and_support(z::AbstractVector{<:Real})
+function simplex_projection_and_support(z::AbstractVector)
     d = length(z)
     z_sorted = sort(z; rev=true)
     z_sorted_cumsum = cumsum(z_sorted)
@@ -37,7 +37,7 @@ function simplex_projection_and_support(z::AbstractVector{<:Real})
     return p, s
 end
 
-function ChainRulesCore.rrule(::typeof(sparse_argmax), z::AbstractVector{<:Real})
+function ChainRulesCore.rrule(::typeof(sparse_argmax), z::AbstractVector)
     p, s = simplex_projection_and_support(z)
     S = sum(s)
     function sparse_argmax_pullback(dp)
diff --git a/src/spo/spoplus_loss.jl b/src/spo/spoplus_loss.jl
diff --git a/src/ssvm/ssvm_loss.jl b/src/ssvm/ssvm_loss.jl
diff --git a/src/ssvm/zeroone_baseloss.jl b/src/ssvm/zeroone_baseloss.jl
diff --git a/src/utils/pushforward.jl b/src/utils/pushforward.jl
diff --git a/test/InferOptTestUtils/error.jl b/test/InferOptTestUtils/error.jl
