From cf146af7625c531ccc1f23b1486cc73a7bf10207 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sat, 19 Jul 2025 15:52:58 +0200 Subject: [PATCH 1/5] Add sizes computation --- src/ArrayDiff.jl | 1 + src/reverse_mode.jl | 43 ++++++++++------ src/sizes.jl | 123 ++++++++++++++++++++++++++++++++++++++++++++ src/types.jl | 5 +- test/ArrayDiff.jl | 71 +++++++++++++++++++++++++ test/runtests.jl | 2 + 6 files changed, 227 insertions(+), 18 deletions(-) create mode 100644 src/sizes.jl create mode 100644 test/ArrayDiff.jl diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl index 1224763..1f08c1e 100644 --- a/src/ArrayDiff.jl +++ b/src/ArrayDiff.jl @@ -50,6 +50,7 @@ import NaNMath: include("Coloring/Coloring.jl") include("graph_tools.jl") +include("sizes.jl") include("types.jl") include("utils.jl") diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl index 84cc358..4f54b77 100644 --- a/src/reverse_mode.jl +++ b/src/reverse_mode.jl @@ -109,20 +109,22 @@ function _forward_eval( # children, so a backwards pass through f.nodes is a forward pass through # the tree. children_arr = SparseArrays.rowvals(f.adj) + fill!(f.partials_storage, zero(T)) for k in length(f.nodes):-1:1 node = f.nodes[k] - f.partials_storage[k] = zero(T) + # Storage index if scalar + j = last(_storage_range(f.sizes, k)) if node.type == Nonlinear.NODE_VARIABLE - f.forward_storage[k] = x[node.index] + f.forward_storage[j] = x[node.index] # This should never happen, because we will have replaced these by now. # elseif node.type == Nonlinear.NODE_MOI_VARIABLE # f.forward_storage[k] = x[node.index] elseif node.type == Nonlinear.NODE_VALUE - f.forward_storage[k] = f.const_values[node.index] + f.forward_storage[j] = f.const_values[node.index] elseif node.type == Nonlinear.NODE_SUBEXPRESSION - f.forward_storage[k] = d.subexpression_forward_values[node.index] + f.forward_storage[j] = d.subexpression_forward_values[node.index] elseif node.type == Nonlinear.NODE_PARAMETER - f.forward_storage[k] = d.data.parameters[node.index] + f.forward_storage[j] = d.data.parameters[node.index] elseif node.type == Nonlinear.NODE_CALL_MULTIVARIATE children_indices = SparseArrays.nzrange(f.adj, k) N = length(children_indices) @@ -130,23 +132,29 @@ function _forward_eval( # With appropriate benchmarking, the special-cased if-statements can # be removed in favor of the generic user-defined function case. if node.index == 1 # :+ - tmp_sum = zero(T) - for c_idx in children_indices - @inbounds ix = children_arr[c_idx] - @inbounds f.partials_storage[ix] = one(T) - @inbounds tmp_sum += f.forward_storage[ix] + for j in _eachindex(f.sizes, k) + tmp_sum = zero(T) + for c_idx in children_indices + ix = children_arr[c_idx] + _setindex!(f.partials_storage, one(T), f.sizes, ix, j) + f.partials_storage[ix] = one(T) + tmp_sum += _getindex(f.forward_storage, f.sizes, ix, j) + end + _setindex!(f.forward_storage, tmp_sum, f.sizes, k, j) end - f.forward_storage[k] = tmp_sum elseif node.index == 2 # :- @assert N == 2 child1 = first(children_indices) @inbounds ix1 = children_arr[child1] @inbounds ix2 = children_arr[child1+1] - @inbounds tmp_sub = f.forward_storage[ix1] - @inbounds tmp_sub -= f.forward_storage[ix2] - @inbounds f.partials_storage[ix1] = one(T) - @inbounds f.partials_storage[ix2] = -one(T) - f.forward_storage[k] = tmp_sub + for j in _eachindex(f.sizes, k) + tmp_sub = _getindex(f.forward_storage, f.sizes, ix1, j) + tmp_sub -= _getindex(f.forward_storage, f.sizes, ix2, j) + tmp_sub = _getindex(f.partials_storage, f.sizes, ix1, one(T)) + _setindex!(f.partials_storage, one(T), f.sizes, ix1, j) + _setindex!(f.partials_storage, -one(T), f.sizes, ix2, j) + _setindex!(f.forward_storage, tmp_sub, f.sizes, k, j) + end elseif node.index == 3 # :* tmp_prod = one(T) for c_idx in children_indices @@ -221,7 +229,8 @@ function _forward_eval( @inbounds f.partials_storage[children_arr[idx1+2]] = !(condition == 1) f.forward_storage[k] = ifelse(condition == 1, lhs, rhs) - else + else # atan, min, max or vect + @show node.index f_input = _UnsafeVectorView(d.jac_storage, N) ∇f = _UnsafeVectorView(d.user_output_buffer, N) for (r, i) in enumerate(children_indices) diff --git a/src/sizes.jl b/src/sizes.jl new file mode 100644 index 0000000..0274a58 --- /dev/null +++ b/src/sizes.jl @@ -0,0 +1,123 @@ +""" + struct Sizes + ndims::Vector{Int} + size_offset::Vector{Int} + size::Vector{Int} + storage_offset::Vector{Int} + end + +The node at index `k` is an array of `ndims[k]` dimensions and size `sizes[size_offset[k] .+ (1:ndims[k])]`. +Note that `size_offset` is a nonincreasing vector so that `sizes` can be filled in a forward pass, +which goes through the nodes in decreasing index order. +""" +struct Sizes + ndims::Vector{Int} + size_offset::Vector{Int} + size::Vector{Int} + storage_offset::Vector{Int} +end + +function _size(sizes::Sizes, k::Int, dim::Int) + return sizes.size[sizes.size_offset[k] + dim] +end + +function _size(sizes::Sizes, k::Int) + return view(sizes.size, sizes.size_offset[k] .+ Base.OneTo(dim)) +end + +function _length(sizes::Sizes, k::Int) + if sizes.ndims[k] == 0 + return 1 + else + return prod(_size(sizes, k)) + end +end + +_eachindex(sizes::Sizes, k) = Base.OneTo(_length(sizes, k)) + +_length(sizes::Sizes) = sizes.storage_offset[end] + +function _storage_range(sizes::Sizes, k::Int) + return sizes.storage_offset[k] .+ _eachindex(sizes, k) +end + +# /!\ Can only be called in decreasing `k` order +function _add_size!(sizes::Sizes, k::Int, size::Tuple) + sizes.ndims[k] = length(size) + sizes.size_offset[k] = length(sizes.size) + append!(sizes.size, size) + return +end + +function _copy_size!(sizes::Sizes, k::Int, child::Int) + sizes.ndims[k] = sizes.ndims[child] + sizes.size_offset[k] = length(sizes.size) + for i in (sizes.size_offset[child] .+ Base.OneTo(sizes.ndims[child])) + push!(sizes.size, sizes.size[i]) + end + return +end + +function _assert_scalar_children(sizes, children_arr, children_indices, op) + for c_idx in children_indices + @inbounds ix = children_arr[c_idx] + # We don't support nested vectors of vectors, + # we only support real numbers and array of real numbers + @assert sizes.ndims[ix] == 0 "Array argument when expected scalar argument for operator `$op`" + end +end + +function _infer_sizes( + nodes::Vector{Nonlinear.Node}, + adj::SparseArrays.SparseMatrixCSC{Bool,Int}, +) + sizes = Sizes(zeros(Int, length(nodes)), Int[], zeros(Int, length(nodes)), zeros(Int, length(nodes) + 1)) + operators = d.data.operators + children_arr = SparseArrays.rowvals(adj) + for k in length(nodes):-1:1 + node = nodes[k] + children_indices = SparseArrays.nzrange(adj, k) + N = length(children_indices) + if node.type == Nonlinear.NODE_SUBEXPRESSION + error("Subexpressions not supported yet") + elseif node.type == Nonlinear.NODE_CALL_MULTIVARIATE + op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index] + if op == :vect + _assert_scalar_children(sizes, children_arr, children_indices, op) + _add_size!(sizes, k, (N,)) + elseif op == :+ || op == :- + # TODO assert all arguments have same size + _copy_size!(sizes, k, children_arr[first(children_indices)]) + elseif op == :* + # TODO assert compatible sizes and all ndims should be 0 or 2 + first_matrix = findfirst(children_indices) do i + !iszero(sizes.ndims[children_arr[i]]) + end + if !isnothing(first_matrix) + last_matrix = findfirst(children_indices) do i + !iszero(sizes.ndims[children_arr[i]]) + end + _add_size!(sizes, k, (_size(sizes, first_matrix, 1), _size(sizes, last_matrix, sizes.ndims[last_matrix]))) + end + elseif op == :^ || op == :/ + @assert N == 2 + _assert_scalar_children(sizes, children_arr, children_indices[2:end], op) + _copy_size!(sizes, k, children_arr[first(children_indices)]) + else + _assert_scalar_children(sizes, children_arr, children_indices, op) + end + elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE + @assert N == 1 + op = MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS[node.index] + if op == :+ || op == :- + _copy_size!(sizes, k, children_arr[first(children_indices)]) + else + _assert_scalar_children(sizes, children_arr, children_indices, op) + end + end + end + for k in eachindex(nodes) + sizes.storage_offset[k + 1] = sizes.storage_offset[k] + _length(sizes, k) + end + return sizes +end diff --git a/src/types.jl b/src/types.jl index 5dad1b8..4f1f19a 100644 --- a/src/types.jl +++ b/src/types.jl @@ -7,6 +7,7 @@ struct _SubexpressionStorage nodes::Vector{Nonlinear.Node} adj::SparseArrays.SparseMatrixCSC{Bool,Int} + sizes::Sizes const_values::Vector{Float64} forward_storage::Vector{Float64} partials_storage::Vector{Float64} @@ -21,10 +22,12 @@ struct _SubexpressionStorage partials_storage_ϵ::Vector{Float64}, linearity::Linearity, ) - N = length(nodes) + sizes = _infer_sizes(nodes, adj) + N = _length(sizes) return new( nodes, adj, + _infer_sizes(nodes, adj), const_values, zeros(N), # forward_storage, zeros(N), # partials_storage, diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl new file mode 100644 index 0000000..99b7be3 --- /dev/null +++ b/test/ArrayDiff.jl @@ -0,0 +1,71 @@ +module TestArrayDiff + +using Test +import LinearAlgebra +import SparseArrays + +import MathOptInterface as MOI +const Nonlinear = MOI.Nonlinear + +import ArrayDiff +const Coloring = ArrayDiff.Coloring + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$(name)", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end + return +end + +function test_objective_quadratic_univariate() + x = MOI.VariableIndex(1) + scalar = Nonlinear.Model() + Nonlinear.set_objective(model, :($x * $x)) + vector = Nonlinear.Model() + Nonlinear.set_objective(vector, :([$x] * [$x])) + evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x]) + MOI.initialize(evaluator, [:Grad, :Jac, :Hess]) + @test MOI.eval_objective(evaluator, [1.2]) == 1.2^2 + 1 + g = [NaN] + MOI.eval_objective_gradient(evaluator, g, [1.2]) + @test g == [2.4] + @test MOI.hessian_objective_structure(evaluator) == [(1, 1)] + H = [NaN] + MOI.eval_hessian_objective(evaluator, H, [1.2]) + @test H == [2.0] + @test MOI.hessian_lagrangian_structure(evaluator) == [(1, 1)] + H = [NaN] + MOI.eval_hessian_lagrangian(evaluator, H, [1.2], 1.5, Float64[]) + @test H == 1.5 .* [2.0] + MOI.eval_hessian_lagrangian_product( + evaluator, + H, + [1.2], + [1.2], + 1.5, + Float64[], + ) + @test H == [1.5 * 2.0 * 1.2] + return +end + +end # module + +TestArrayDiff.runtests() + +import MathOptInterface as MOI +import ArrayDiff +using Test +const Nonlinear = MOI.Nonlinear +model = Nonlinear.Model() +x = MOI.VariableIndex(1) +Nonlinear.set_objective(model, :(dot([$x], [$x]))) +evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x]) +MOI.initialize(evaluator, [:Grad]) +@test MOI.eval_objective(evaluator, [1.2]) == 1.2^2 +g = [NaN] +MOI.eval_objective_gradient(evaluator, g, [1.2]) diff --git a/test/runtests.jl b/test/runtests.jl index ab66e5e..e290998 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1 +1,3 @@ include("ReverseAD.jl") + +import SparseArrays From d6b90e593fa3622b61e17c76ae8d92550e0e24fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sat, 19 Jul 2025 16:01:36 +0200 Subject: [PATCH 2/5] Fixes --- src/sizes.jl | 7 ++++--- test/ArrayDiff.jl | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/sizes.jl b/src/sizes.jl index 0274a58..e97c10d 100644 --- a/src/sizes.jl +++ b/src/sizes.jl @@ -22,7 +22,7 @@ function _size(sizes::Sizes, k::Int, dim::Int) end function _size(sizes::Sizes, k::Int) - return view(sizes.size, sizes.size_offset[k] .+ Base.OneTo(dim)) + return view(sizes.size, sizes.size_offset[k] .+ Base.OneTo(sizes.ndims[k])) end function _length(sizes::Sizes, k::Int) @@ -71,8 +71,7 @@ function _infer_sizes( nodes::Vector{Nonlinear.Node}, adj::SparseArrays.SparseMatrixCSC{Bool,Int}, ) - sizes = Sizes(zeros(Int, length(nodes)), Int[], zeros(Int, length(nodes)), zeros(Int, length(nodes) + 1)) - operators = d.data.operators + sizes = Sizes(zeros(Int, length(nodes)), zeros(Int, length(nodes)), Int[], zeros(Int, length(nodes) + 1)) children_arr = SparseArrays.rowvals(adj) for k in length(nodes):-1:1 node = nodes[k] @@ -85,6 +84,8 @@ function _infer_sizes( if op == :vect _assert_scalar_children(sizes, children_arr, children_indices, op) _add_size!(sizes, k, (N,)) + elseif op == :dot + # TODO assert all arguments have same size elseif op == :+ || op == :- # TODO assert all arguments have same size _copy_size!(sizes, k, children_arr[first(children_indices)]) diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl index 99b7be3..761237c 100644 --- a/test/ArrayDiff.jl +++ b/test/ArrayDiff.jl @@ -66,6 +66,21 @@ x = MOI.VariableIndex(1) Nonlinear.set_objective(model, :(dot([$x], [$x]))) evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x]) MOI.initialize(evaluator, [:Grad]) +sizes = evaluator.backend.objective.expr.sizes @test MOI.eval_objective(evaluator, [1.2]) == 1.2^2 +@test sizes.ndims == [0, 1, 0, 1, 0] +@test sizes.size_offset == [0, 1, 0, 0, 0] +@test sizes.size == [1, 1] +@test sizes.storage_offset == [0, 1, 2, 3, 4, 5] + +y = MOI.VariableIndex(1) +Nonlinear.set_objective(model, :(dot([$x, $y] - [1, 2], -[1, 2] + [$x, $y]))) +MOI.initialize(evaluator, [:Grad]) +sizes = evaluator.backend.objective.expr.sizes +@test MOI.eval_objective(evaluator, [1.2]) == 1.2^2 +@test sizes.ndims == [0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0] +@test sizes.size_offset == [0, 6, 5, 0, 0, 4, 0, 0, 3, 2, 1, 0, 0, 0, 0, 0] +@test sizes.size == [2, 2, 2, 2, 2, 2, 2] +@test sizes.storage_offset == [0, 1, 3, 5, 6, 7, 9, 10, 11, 13, 15, 17, 18, 19, 21, 22, 23] g = [NaN] MOI.eval_objective_gradient(evaluator, g, [1.2]) From c90c007628cead8b4832fb22b4d7ed8a3676dabc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sat, 19 Jul 2025 16:06:18 +0200 Subject: [PATCH 3/5] Fixes --- src/sizes.jl | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/sizes.jl b/src/sizes.jl index e97c10d..701426c 100644 --- a/src/sizes.jl +++ b/src/sizes.jl @@ -41,6 +41,14 @@ function _storage_range(sizes::Sizes, k::Int) return sizes.storage_offset[k] .+ _eachindex(sizes, k) end +function _getindex(x, sizes::Sizes, k::Int, j) + return x[sizes.storage_offset[k] + j] +end + +function _setindex!(x, value, sizes::Sizes, k::Int, j) + return x[sizes.storage_offset[k] + j] = value +end + # /!\ Can only be called in decreasing `k` order function _add_size!(sizes::Sizes, k::Int, size::Tuple) sizes.ndims[k] = length(size) @@ -77,9 +85,7 @@ function _infer_sizes( node = nodes[k] children_indices = SparseArrays.nzrange(adj, k) N = length(children_indices) - if node.type == Nonlinear.NODE_SUBEXPRESSION - error("Subexpressions not supported yet") - elseif node.type == Nonlinear.NODE_CALL_MULTIVARIATE + if node.type == Nonlinear.NODE_CALL_MULTIVARIATE op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index] if op == :vect _assert_scalar_children(sizes, children_arr, children_indices, op) From 3dcf8cc8ddf1f108d1af0e33d7cf1f5905e64271 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sat, 19 Jul 2025 16:10:30 +0200 Subject: [PATCH 4/5] Fix tests --- src/reverse_mode.jl | 2 -- src/sizes.jl | 8 ++++++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl index 4f54b77..faa14a5 100644 --- a/src/reverse_mode.jl +++ b/src/reverse_mode.jl @@ -150,7 +150,6 @@ function _forward_eval( for j in _eachindex(f.sizes, k) tmp_sub = _getindex(f.forward_storage, f.sizes, ix1, j) tmp_sub -= _getindex(f.forward_storage, f.sizes, ix2, j) - tmp_sub = _getindex(f.partials_storage, f.sizes, ix1, one(T)) _setindex!(f.partials_storage, one(T), f.sizes, ix1, j) _setindex!(f.partials_storage, -one(T), f.sizes, ix2, j) _setindex!(f.forward_storage, tmp_sub, f.sizes, k, j) @@ -230,7 +229,6 @@ function _forward_eval( !(condition == 1) f.forward_storage[k] = ifelse(condition == 1, lhs, rhs) else # atan, min, max or vect - @show node.index f_input = _UnsafeVectorView(d.jac_storage, N) ∇f = _UnsafeVectorView(d.user_output_buffer, N) for (r, i) in enumerate(children_indices) diff --git a/src/sizes.jl b/src/sizes.jl index 701426c..fc78520 100644 --- a/src/sizes.jl +++ b/src/sizes.jl @@ -86,6 +86,10 @@ function _infer_sizes( children_indices = SparseArrays.nzrange(adj, k) N = length(children_indices) if node.type == Nonlinear.NODE_CALL_MULTIVARIATE + if !(node.index in eachindex(MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS)) + # TODO user-defined operators + continue + end op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index] if op == :vect _assert_scalar_children(sizes, children_arr, children_indices, op) @@ -114,6 +118,10 @@ function _infer_sizes( _assert_scalar_children(sizes, children_arr, children_indices, op) end elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE + if !(node.index in eachindex(MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS)) + # TODO user-defined operators + continue + end @assert N == 1 op = MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS[node.index] if op == :+ || op == :- From 915d36fb2c0430208ba71c087ade06a7986b9d93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sat, 19 Jul 2025 16:10:36 +0200 Subject: [PATCH 5/5] Fix format --- src/sizes.jl | 66 +++++++++++++++++++++++++++++++++++++---------- test/ArrayDiff.jl | 3 ++- 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/src/sizes.jl b/src/sizes.jl index fc78520..b745156 100644 --- a/src/sizes.jl +++ b/src/sizes.jl @@ -18,7 +18,7 @@ struct Sizes end function _size(sizes::Sizes, k::Int, dim::Int) - return sizes.size[sizes.size_offset[k] + dim] + return sizes.size[sizes.size_offset[k]+dim] end function _size(sizes::Sizes, k::Int) @@ -42,11 +42,11 @@ function _storage_range(sizes::Sizes, k::Int) end function _getindex(x, sizes::Sizes, k::Int, j) - return x[sizes.storage_offset[k] + j] + return x[sizes.storage_offset[k]+j] end function _setindex!(x, value, sizes::Sizes, k::Int, j) - return x[sizes.storage_offset[k] + j] = value + return x[sizes.storage_offset[k]+j] = value end # /!\ Can only be called in decreasing `k` order @@ -79,20 +79,33 @@ function _infer_sizes( nodes::Vector{Nonlinear.Node}, adj::SparseArrays.SparseMatrixCSC{Bool,Int}, ) - sizes = Sizes(zeros(Int, length(nodes)), zeros(Int, length(nodes)), Int[], zeros(Int, length(nodes) + 1)) + sizes = Sizes( + zeros(Int, length(nodes)), + zeros(Int, length(nodes)), + Int[], + zeros(Int, length(nodes) + 1), + ) children_arr = SparseArrays.rowvals(adj) for k in length(nodes):-1:1 node = nodes[k] children_indices = SparseArrays.nzrange(adj, k) N = length(children_indices) if node.type == Nonlinear.NODE_CALL_MULTIVARIATE - if !(node.index in eachindex(MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS)) + if !( + node.index in + eachindex(MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS) + ) # TODO user-defined operators continue end op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index] if op == :vect - _assert_scalar_children(sizes, children_arr, children_indices, op) + _assert_scalar_children( + sizes, + children_arr, + children_indices, + op, + ) _add_size!(sizes, k, (N,)) elseif op == :dot # TODO assert all arguments have same size @@ -102,23 +115,43 @@ function _infer_sizes( elseif op == :* # TODO assert compatible sizes and all ndims should be 0 or 2 first_matrix = findfirst(children_indices) do i - !iszero(sizes.ndims[children_arr[i]]) + return !iszero(sizes.ndims[children_arr[i]]) end if !isnothing(first_matrix) last_matrix = findfirst(children_indices) do i - !iszero(sizes.ndims[children_arr[i]]) + return !iszero(sizes.ndims[children_arr[i]]) end - _add_size!(sizes, k, (_size(sizes, first_matrix, 1), _size(sizes, last_matrix, sizes.ndims[last_matrix]))) + _add_size!( + sizes, + k, + ( + _size(sizes, first_matrix, 1), + _size(sizes, last_matrix, sizes.ndims[last_matrix]), + ), + ) end elseif op == :^ || op == :/ @assert N == 2 - _assert_scalar_children(sizes, children_arr, children_indices[2:end], op) + _assert_scalar_children( + sizes, + children_arr, + children_indices[2:end], + op, + ) _copy_size!(sizes, k, children_arr[first(children_indices)]) else - _assert_scalar_children(sizes, children_arr, children_indices, op) + _assert_scalar_children( + sizes, + children_arr, + children_indices, + op, + ) end elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE - if !(node.index in eachindex(MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS)) + if !( + node.index in + eachindex(MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS) + ) # TODO user-defined operators continue end @@ -127,12 +160,17 @@ function _infer_sizes( if op == :+ || op == :- _copy_size!(sizes, k, children_arr[first(children_indices)]) else - _assert_scalar_children(sizes, children_arr, children_indices, op) + _assert_scalar_children( + sizes, + children_arr, + children_indices, + op, + ) end end end for k in eachindex(nodes) - sizes.storage_offset[k + 1] = sizes.storage_offset[k] + _length(sizes, k) + sizes.storage_offset[k+1] = sizes.storage_offset[k] + _length(sizes, k) end return sizes end diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl index 761237c..f920107 100644 --- a/test/ArrayDiff.jl +++ b/test/ArrayDiff.jl @@ -81,6 +81,7 @@ sizes = evaluator.backend.objective.expr.sizes @test sizes.ndims == [0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0] @test sizes.size_offset == [0, 6, 5, 0, 0, 4, 0, 0, 3, 2, 1, 0, 0, 0, 0, 0] @test sizes.size == [2, 2, 2, 2, 2, 2, 2] -@test sizes.storage_offset == [0, 1, 3, 5, 6, 7, 9, 10, 11, 13, 15, 17, 18, 19, 21, 22, 23] +@test sizes.storage_offset == + [0, 1, 3, 5, 6, 7, 9, 10, 11, 13, 15, 17, 18, 19, 21, 22, 23] g = [NaN] MOI.eval_objective_gradient(evaluator, g, [1.2])