diff --git a/src/Nonlinear/ReverseAD/utils.jl b/src/Nonlinear/ReverseAD/utils.jl index 3b5a95369b..231eb173b8 100644 --- a/src/Nonlinear/ReverseAD/utils.jl +++ b/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,) diff --git a/test/Nonlinear/ReverseAD.jl b/test/Nonlinear/ReverseAD.jl index 8328d89d1b..43fbb28b4f 100644 --- a/test/Nonlinear/ReverseAD.jl +++ b/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()