Make it work on Float32

Author: blegat

src/ArrayDiff.jl

@@ -20,10 +20,11 @@ Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
 
 The type parameter `S` is the storage type used for the AD tape (forward,
 partials, and reverse storage of each subexpression). It must satisfy
-`S<:AbstractVector{Float64}`. Defaults to `Vector{Float64}`. Pass a different
-`S` (for example `CuVector{Float64}`) to keep the tape on a GPU.
+`S<:AbstractVector{<:Real}`. Defaults to `Vector{Float64}`. Pass a different
+`S` (for example `Vector{Float32}` or `CuVector{Float64}`) to run AD in
+another precision or keep the tape on a GPU.
 """
-struct Mode{S<:AbstractVector{Float64}} <:
+struct Mode{S<:AbstractVector{<:Real}} <:
        MOI.Nonlinear.AbstractAutomaticDifferentiation end
 
 Mode() = Mode{Vector{Float64}}()
@@ -65,7 +66,7 @@ include("evaluator.jl")
 include("array_nonlinear_function.jl")
 include("parse_moi.jl")
 
-model(::Mode{S}) where {S} = Model()
+model(::Mode{S}) where {S} = Model{eltype(S)}()
 
 # Extend MOI.Nonlinear.set_objective so that solvers calling
 # MOI.Nonlinear.set_objective(arraydiff_model, snf) dispatch here.
@@ -84,8 +85,8 @@ function Evaluator(
     model::ArrayDiff.Model,
     ::Mode{S},
     ordered_variables::Vector{MOI.VariableIndex},
-) where {S<:AbstractVector{Float64}}
-    return Evaluator(model, NLPEvaluator{S}(model, ordered_variables))
+) where {S<:AbstractVector{<:Real}}
+    return Evaluator(model, NLPEvaluator{eltype(S),S}(model, ordered_variables))
 end
 
 # Called by solvers via MOI.Nonlinear.Evaluator(nlp_model, ad_backend, vars).

src/JuMP/moi_bridge.jl

@@ -12,9 +12,9 @@ function _to_moi_arg(x::GenericArrayExpr{V,N}) where {V,N}
     return ArrayNonlinearFunction{N}(x.head, args, x.size, x.broadcasted)
 end
 
-_to_moi_arg(x::Matrix{Float64}) = x
+_to_moi_arg(x::Matrix{<:Real}) = x
 
-_to_moi_arg(x::Real) = Float64(x)
+_to_moi_arg(x::Real) = x
 
 function JuMP.moi_function(x::GenericArrayExpr{V,N}) where {V,N}
     return _to_moi_arg(x)

src/array_nonlinear_function.jl

@@ -87,7 +87,7 @@ function _map_indices_arg(index_map::F, x::ArrayOfContiguousVariables) where {F}
     return MOI.Utilities.map_indices(index_map, x)
 end
 
-function _map_indices_arg(::F, x::Matrix{Float64}) where {F}
+function _map_indices_arg(::F, x::Matrix{<:Real}) where {F}
     return x
 end
 

src/mathoptinterface_api.jl

@@ -20,9 +20,9 @@ function MOI.features_available(d::NLPEvaluator)
 end
 
 function MOI.initialize(
-    d::NLPEvaluator{S},
+    d::NLPEvaluator{T,S},
     requested_features::Vector{Symbol},
-) where {S<:AbstractVector{Float64}}
+) where {T<:Real,S<:AbstractVector{T}}
     # Check that we support the features requested by the user.
     available_features = MOI.features_available(d)
     for feature in requested_features
@@ -40,10 +40,10 @@ function MOI.initialize(
     end
     d.objective = nothing
     d.residual = nothing
-    d.user_output_buffer = zeros(largest_user_input_dimension)
-    d.jac_storage = zeros(max(N, largest_user_input_dimension))
-    d.constraints = _FunctionStorage{S}[]
-    d.last_x = fill(NaN, N)
+    d.user_output_buffer = zeros(T, largest_user_input_dimension)
+    d.jac_storage = zeros(T, max(N, largest_user_input_dimension))
+    d.constraints = _FunctionStorage{T,S}[]
+    d.last_x = fill(T(NaN), N)
     d.want_hess = :Hess in requested_features
     want_hess_storage = (:HessVec in requested_features) || d.want_hess
     coloring_storage = Coloring.IndexedSet(N)
@@ -67,9 +67,9 @@ function MOI.initialize(
     subexpression_edgelist =
         Vector{Set{Tuple{Int,Int}}}(undef, num_subexpressions)
     d.subexpressions =
-        Vector{_SubexpressionStorage{S}}(undef, num_subexpressions)
-    d.subexpression_forward_values = zeros(num_subexpressions)
-    d.subexpression_reverse_values = zeros(num_subexpressions)
+        Vector{_SubexpressionStorage{T,S}}(undef, num_subexpressions)
+    d.subexpression_forward_values = zeros(T, num_subexpressions)
+    d.subexpression_reverse_values = zeros(T, num_subexpressions)
     for k in d.subexpression_order
         # Only load expressions which actually are used
         d.subexpression_forward_values[k] = NaN
@@ -145,6 +145,7 @@ function MOI.initialize(
             moi_index_to_consecutive_index,
             shared_partials_storage_ϵ,
             d,
+            S,
         )
         residual = _FunctionStorage(
             subexpr,
@@ -235,7 +236,7 @@ function MOI.eval_objective_gradient(d::NLPEvaluator, g, x)
         error("No nonlinear objective.")
     end
     _reverse_mode(d, x)
-    fill!(g, 0.0)
+    fill!(g, zero(eltype(g)))
     _extract_reverse_pass(g, d, something(d.objective))
     return
 end

src/reverse_mode.jl

@@ -945,9 +945,9 @@ function _extract_reverse_pass(
     f::_FunctionStorage,
 ) where {T}
     for i in f.dependent_subexpressions
-        d.subexpression_reverse_values[i] = 0.0
+        d.subexpression_reverse_values[i] = zero(T)
     end
-    _extract_reverse_pass_inner(g, f, d.subexpression_reverse_values, 1.0)
+    _extract_reverse_pass_inner(g, f, d.subexpression_reverse_values, one(T))
     for i in length(f.dependent_subexpressions):-1:1
         k = f.dependent_subexpressions[i]
         _extract_reverse_pass_inner(

src/sizes.jl

@@ -438,11 +438,11 @@ function _infer_sizes(
     return sizes
 end
 
-struct _SubexpressionStorage{S<:AbstractVector{Float64}}
+struct _SubexpressionStorage{T<:Real,S<:AbstractVector{T}}
     nodes::Vector{Node}
     adj::SparseArrays.SparseMatrixCSC{Bool,Int}
     sizes::Sizes
-    const_values::Vector{Float64}
+    const_values::Vector{T}
     forward_storage::S
     partials_storage::S
     reverse_storage::S
@@ -452,19 +452,19 @@ struct _SubexpressionStorage{S<:AbstractVector{Float64}}
     function _SubexpressionStorage(
         nodes::Vector{Node},
         adj::SparseArrays.SparseMatrixCSC{Bool,Int},
-        const_values::Vector{Float64},
+        const_values::Vector{T},
         block_shapes::Dict{Int,Vector{Int}},
         partials_storage_ϵ::Vector{Float64},
         linearity::Linearity,
-        ::Type{S} = Vector{Float64},
-    ) where {S<:AbstractVector{Float64}}
+        ::Type{S} = Vector{T},
+    ) where {T<:Real,S<:AbstractVector{T}}
         sizes = _infer_sizes(nodes, adj, block_shapes)
         N = _length(sizes)
         # Pre-load value blocks into forward_storage once at construction;
         # each block is a contiguous-to-contiguous bulk copy. Individual
         # `NODE_VALUE` scalars (rare — exponents, constant divisors, etc) and
         # variable nodes are loaded by `_forward_eval` in the per-node loop.
-        cpu_buffer = zeros(N)
+        cpu_buffer = zeros(T, N)
         for k in 1:length(nodes)
             node = nodes[k]
             if node.type == NODE_VALUE_BLOCK
@@ -475,14 +475,14 @@ struct _SubexpressionStorage{S<:AbstractVector{Float64}}
             end
         end
         forward_storage = convert(S, cpu_buffer)
-        return new{S}(
+        return new{T,S}(
             nodes,
             adj,
             sizes,
             const_values,
             forward_storage,
-            fill!(S(undef, N), 0.0),  # partials_storage,
-            fill!(S(undef, N), 0.0),  # reverse_storage,
+            fill!(S(undef, N), zero(T)),  # partials_storage,
+            fill!(S(undef, N), zero(T)),  # reverse_storage,
             partials_storage_ϵ,
             linearity,
         )

src/types.jl

@@ -103,7 +103,7 @@ function _subexpression_and_linearity(
     partials_storage_ϵ::Vector{Float64},
     d,
     ::Type{S} = Vector{Float64},
-) where {S<:AbstractVector{Float64}}
+) where {S<:AbstractVector{<:Real}}
     nodes = _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index)
     adj = adjacency_matrix(nodes)
     linearity = if d.want_hess
@@ -114,7 +114,7 @@ function _subexpression_and_linearity(
     return _SubexpressionStorage(
         nodes,
         adj,
-        convert(Vector{Float64}, expr.values),
+        convert(Vector{eltype(S)}, expr.values),
         copy(expr.block_shapes),
         partials_storage_ϵ,
         linearity[1],
@@ -123,28 +123,28 @@ function _subexpression_and_linearity(
     linearity
 end
 
-struct _FunctionStorage{S<:AbstractVector{Float64}}
-    expr::_SubexpressionStorage{S}
+struct _FunctionStorage{T<:Real,S<:AbstractVector{T}}
+    expr::_SubexpressionStorage{T,S}
     grad_sparsity::Vector{Int}
     # Nonzero pattern of Hessian matrix
     hess_I::Vector{Int}
     hess_J::Vector{Int}
     rinfo::Coloring.RecoveryInfo # coloring info for hessians
-    seed_matrix::Matrix{Float64}
+    seed_matrix::Matrix{T}
     # subexpressions which this function depends on, ordered for forward pass.
     dependent_subexpressions::Vector{Int}
 
     function _FunctionStorage(
-        expr::_SubexpressionStorage{S},
+        expr::_SubexpressionStorage{T,S},
         num_variables,
         coloring_storage::Coloring.IndexedSet,
         want_hess::Bool,
-        subexpressions::Vector{_SubexpressionStorage{S}},
+        subexpressions::Vector{_SubexpressionStorage{T,S}},
         dependent_subexpressions,
         subexpression_edgelist,
         subexpression_variables,
         linearity::Vector{Linearity},
-    ) where {S<:AbstractVector{Float64}}
+    ) where {T<:Real,S<:AbstractVector{T}}
         empty!(coloring_storage)
         _compute_gradient_sparsity!(coloring_storage, expr)
         for k in dependent_subexpressions
@@ -166,7 +166,7 @@ struct _FunctionStorage{S<:AbstractVector{Float64}}
                 coloring_storage,
             )
             seed_matrix = Coloring.seed_matrix(rinfo)
-            return new{S}(
+            return new{T,S}(
                 expr,
                 grad_sparsity,
                 hess_I,
@@ -176,13 +176,13 @@ struct _FunctionStorage{S<:AbstractVector{Float64}}
                 dependent_subexpressions,
             )
         else
-            return new{S}(
+            return new{T,S}(
                 expr,
                 grad_sparsity,
                 Int[],
                 Int[],
                 Coloring.RecoveryInfo(),
-                Array{Float64}(undef, 0, 0),
+                Array{T}(undef, 0, 0),
                 dependent_subexpressions,
             )
         end
@@ -305,30 +305,30 @@ interface.
 !!! warning
     Before using, you must initialize the evaluator using `MOI.initialize`.
 """
-mutable struct NLPEvaluator{S<:AbstractVector{Float64}} <:
+mutable struct NLPEvaluator{T<:Real,S<:AbstractVector{T}} <:
                MOI.AbstractNLPEvaluator
     data::Model
     ordered_variables::Vector{MOI.VariableIndex}
 
-    objective::Union{Nothing,_FunctionStorage{S}}
-    residual::Union{Nothing,_FunctionStorage{S}}
-    constraints::Vector{_FunctionStorage{S}}
-    subexpressions::Vector{_SubexpressionStorage{S}}
+    objective::Union{Nothing,_FunctionStorage{T,S}}
+    residual::Union{Nothing,_FunctionStorage{T,S}}
+    constraints::Vector{_FunctionStorage{T,S}}
+    subexpressions::Vector{_SubexpressionStorage{T,S}}
     subexpression_order::Vector{Int}
     # Storage for the subexpressions in reverse-mode automatic differentiation.
-    subexpression_forward_values::Vector{Float64}
-    subexpression_reverse_values::Vector{Float64}
+    subexpression_forward_values::Vector{T}
+    subexpression_reverse_values::Vector{T}
     subexpression_linearity::Vector{Linearity}
 
     # A cache of the last x. This is used to guide whether we need to re-run
     # reverse-mode automatic differentiation.
-    last_x::Vector{Float64}
+    last_x::Vector{T}
 
     # Temporary storage for computing Jacobians. This is also used as temporary
     # storage for the input of multivariate functions.
-    jac_storage::Vector{Float64}
+    jac_storage::Vector{T}
     # Temporary storage for the gradient of multivariate functions
-    user_output_buffer::Vector{Float64}
+    user_output_buffer::Vector{T}
 
     # storage for computing hessians
     # these Float64 vectors are reinterpreted to hold multiple epsilon components
@@ -343,10 +343,10 @@ mutable struct NLPEvaluator{S<:AbstractVector{Float64}} <:
     hessian_sparsity::Vector{Tuple{Int64,Int64}}
     max_chunk::Int # chunk size for which we've allocated storage
 
-    function NLPEvaluator{S}(
+    function NLPEvaluator{T,S}(
         data::Model,
         ordered_variables::Vector{MOI.VariableIndex},
-    ) where {S<:AbstractVector{Float64}}
-        return new{S}(data, ordered_variables)
+    ) where {T<:Real,S<:AbstractVector{T}}
+        return new{T,S}(data, ordered_variables)
     end
 end