[Nonlinear] fix _UnsafeVectorView with [email protected] (#2708)

Author: odow

src/Nonlinear/ReverseAD/utils.jl

@@ -42,13 +42,25 @@ struct _UnsafeVectorView{T} <: DenseVector{T}
     ptr::Ptr{T}
 end
 
-Base.getindex(x::_UnsafeVectorView, i) = unsafe_load(x.ptr, i + x.offset)
+function Base.getindex(x::_UnsafeVectorView, i::Integer)
+    return unsafe_load(x.ptr, i + x.offset)
+end
+
+Base.getindex(x::_UnsafeVectorView, i::CartesianIndex{1}) = getindex(x, i[1])
 
-function Base.setindex!(x::_UnsafeVectorView, value, i)
+function Base.setindex!(x::_UnsafeVectorView, value, i::Integer)
+    # We don't need to worry about `value` being the right type here because
+    # x.ptr is a `::Ptr{T}`, so even though it is called `unsafe_store!`, there
+    # is still a type convertion that happens so that we're not just chucking
+    # the bits of value into `x.ptr`.
     unsafe_store!(x.ptr, value, i + x.offset)
     return value
 end
 
+function Base.setindex!(x::_UnsafeVectorView, value, i::CartesianIndex{1})
+    return setindex!(x, value, i[1])
+end
+
 Base.length(v::_UnsafeVectorView) = v.len
 
 Base.size(v::_UnsafeVectorView) = (v.len,)

test/Nonlinear/ReverseAD.jl

@@ -1304,6 +1304,65 @@ function test_toposort_subexpressions()
     return
 end
 
+function test_eval_user_defined_operator_ForwardDiff_gradient!()
+    model = MOI.Nonlinear.Model()
+    x = MOI.VariableIndex.(1:4)
+    p = MOI.Nonlinear.add_parameter(model, 2.0)
+    ex = MOI.Nonlinear.add_expression(model, :($p * $(x[1])))
+    ψ(x) = sin(x)
+    t(x, y) = x + 3y
+    MOI.Nonlinear.register_operator(model, :ψ, 1, ψ)
+    MOI.Nonlinear.register_operator(model, :t, 2, t)
+    MOI.Nonlinear.add_constraint(
+        model,
+        :($ex^3 + sin($(x[2])) / ψ($(x[2])) + t($(x[3]), $(x[4]))),
+        MOI.LessThan(0.0),
+    )
+    d = MOI.Nonlinear.Evaluator(model, MOI.Nonlinear.SparseReverseMode(), x)
+    MOI.initialize(d, [:Jac])
+    X = [1.1, 1.2, 1.3, 1.4]
+    g = [NaN]
+    MOI.eval_constraint(d, g, X)
+    @test only(g) ≈ 17.148
+    @test MOI.jacobian_structure(d) == [(1, 1), (1, 2), (1, 3), (1, 4)]
+    J = [NaN, NaN, NaN, NaN]
+    MOI.eval_constraint_jacobian(d, J, X)
+    @test J ≈ [2.0^3 * 3.0 * 1.1^2, 0.0, 1.0, 3.0]
+    return
+end
+
+function test_eval_user_defined_operator_type_mismatch()
+    model = MOI.Nonlinear.Model()
+    x = MOI.VariableIndex.(1:4)
+    p = MOI.Nonlinear.add_parameter(model, 2.0)
+    ex = MOI.Nonlinear.add_expression(model, :($p * $(x[1])))
+    ψ(x) = sin(x)
+    t(x, y) = x + 3y
+    function ∇t(ret, x, y)
+        ret[1] = 1      # These are intentionally the wrong type
+        ret[2] = 3 // 1 # These are intentionally the wrong type
+        return
+    end
+    MOI.Nonlinear.register_operator(model, :ψ, 1, ψ, cos)
+    MOI.Nonlinear.register_operator(model, :t, 2, t, ∇t)
+    MOI.Nonlinear.add_constraint(
+        model,
+        :($ex^3 + sin($(x[2])) / ψ($(x[2])) + t($(x[3]), $(x[4]))),
+        MOI.LessThan(0.0),
+    )
+    d = MOI.Nonlinear.Evaluator(model, MOI.Nonlinear.SparseReverseMode(), x)
+    MOI.initialize(d, [:Jac])
+    X = [1.1, 1.2, 1.3, 1.4]
+    g = [NaN]
+    MOI.eval_constraint(d, g, X)
+    @test only(g) ≈ 17.148
+    @test MOI.jacobian_structure(d) == [(1, 1), (1, 2), (1, 3), (1, 4)]
+    J = [NaN, NaN, NaN, NaN]
+    MOI.eval_constraint_jacobian(d, J, X)
+    @test J ≈ [2.0^3 * 3.0 * 1.1^2, 0.0, 1.0, 3.0]
+    return
+end
+
 end  # module
 
 TestReverseAD.runtests()