Fix reshape

src/SymbolicNeuralNetworks.jl

@@ -20,6 +20,8 @@ module SymbolicNeuralNetworks
     export symbolize
     include("utils/symbolize.jl")
 
+    include("utils/create_array.jl")
+
     export AbstractSymbolicNeuralNetwork
     export SymbolicNeuralNetwork, SymbolicModel
     export HamiltonianSymbolicNeuralNetwork, HNNLoss

src/pullback.jl

@@ -93,7 +93,7 @@ function SymbolicPullback(nn::SymbolicNeuralNetwork, loss::NetworkLoss)
     @variables soutput[1:output_dimension(nn.model)]
     symbolic_loss = loss(nn.model, nn.params, nn.input, soutput)
     symbolic_pullbacks = symbolic_pullback(symbolic_loss, nn)
-    pbs_executable = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput)
+    pbs_executable = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput; reduce = +)
     function pbs(input, output, params)
         pullback(::Union{Real, AbstractArray{<:Real}}) = _get_contents(_get_params(pbs_executable(input, output, params)))
         pullback

src/utils/build_function.jl

@@ -22,12 +22,19 @@ function build_nn_function(eq::EqT, nn::AbstractSymbolicNeuralNetwork)
     build_nn_function(eq, nn.params, nn.input)
 end
 
-function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr)
+function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr; reduce = hcat)
     gen_fun = _build_nn_function(eq, sparams, sinput)
-    gen_fun_returned(x, ps) = mapreduce(k -> gen_fun(x, ps, k), hcat, axes(x, 2))
-    gen_fun_returned(x::Union{AbstractVector, Symbolics.Arr}, ps) = gen_fun_returned(reshape(x, length(x), 1), ps)
+    gen_fun_returned(x, ps) = mapreduce(k -> gen_fun(x, ps, k), reduce, axes(x, 2))
+    function gen_fun_returned(x::Union{AbstractVector, Symbolics.Arr}, ps) 
+        output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), ps)
+        # for vectors we do not reshape, as the output may be a matrix
+        output_not_reshaped
+    end
     # check this! (definitely not correct in all cases!)
-    gen_fun_returned(x::AbstractArray{<:Number, 3}, ps) = reshape(gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), ps), size(x, 1), size(x, 2), size(x, 3))
+    function gen_fun_returned(x::AbstractArray{<:Number, 3}, ps) 
+        output_not_reshaped = gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), ps)
+        reshape(output_not_reshaped, size(output_not_reshaped, 1), size(x, 2), size(x, 3))
+    end
     gen_fun_returned
 end
 

src/utils/build_function2.jl

@@ -16,14 +16,31 @@
     build_nn_function(eqs, nn.params, nn.input, soutput)
 end
 
-function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr)
+function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr; reduce = hcat)
     gen_fun = _build_nn_function(eq, sparams, sinput, soutput)
-    gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), +, axes(input, 2))
-    gen_fun_returned(input::AT, output::AT, ps) where {AT <: Union{AbstractVector, Symbolics.Arr}} = gen_fun_returned(reshape(input, length(input), 1), reshape(output, length(output), 1), ps)
-    gen_fun_returned(input::AT, output::AT, ps) where {T, AT <: AbstractArray{T, 3}} = gen_fun_returned(reshape(input, size(input, 1), size(input, 2) * size(input, 3)), reshape(output, size(output, 1), size(output, 2) * size(output, 3)), ps)
+    gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), reduce, axes(input, 2))
+    function gen_fun_returned(x::AT, y::AT, ps) where {AT <: Union{AbstractVector, Symbolics.Arr}}
+        output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), reshape(y, length(y), 1), ps)
+        # for vectors we do not reshape, as the output may be a matrix
+        output_not_reshaped
+    end
+    # check this! (definitely not correct in all cases!)
+    function gen_fun_returned(x::AT, y::AT, ps) where {AT <: AbstractArray{<:Number, 3}} 
+        output_not_reshaped = gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), reshape(y, size(y, 1), size(y, 2) * size(y, 3)), ps)
+        # if arrays are added together then don't reshape!
+        optional_reshape(output_not_reshaped, reduce, x)
+    end
     gen_fun_returned
 end
 
+function optional_reshape(output_not_reshaped::AbstractVecOrMat, ::typeof(+), ::AbstractArray{<:Number, 3})
+    output_not_reshaped
+end
+
+function optional_reshape(output_not_reshaped::AbstractVecOrMat, ::typeof(hcat), input::AbstractArray{<:Number, 3})
+    reshape(output_not_reshaped, size(output_not_reshaped, 1), size(input, 2), size(input, 3))
+end
+
 """
     _build_nn_function(eq, params, sinput, soutput)
 

src/utils/build_function_arrays.jl

@@ -29,8 +29,8 @@ funcs_evaluated = funcs(input, ps)
 (true, true, true)
 ```
 """
-function build_nn_function(eqs::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...)
-    ps_semi = [function_valued_parameters(eq, sparams, sinput...) for eq in eqs]
+function build_nn_function(eqs::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat)
+    ps_semi = [function_valued_parameters(eq, sparams, sinput...; reduce = reduce) for eq in eqs]
 
     _pbs_executable(ps_functions, params, input...) = apply_element_wise(ps_functions, params, input...)
     __pbs_executable(input, params) = _pbs_executable(ps_semi, params, input)
@@ -72,8 +72,8 @@ funcs_evaluated = funcs(input, ps)
 
 Internally this is using [`function_valued_parameters`](@ref) and [`apply_element_wise`](@ref).
 """
-function build_nn_function(eqs::Union{NamedTuple, NeuralNetworkParameters}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...)
-    ps = function_valued_parameters(eqs, sparams, sinput...)
+function build_nn_function(eqs::Union{NamedTuple, NeuralNetworkParameters}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat)
+    ps = function_valued_parameters(eqs, sparams, sinput...; reduce = reduce)
     _pbs_executable(ps::Union{NamedTuple, NeuralNetworkParameters}, params::NeuralNetworkParameters, input::AbstractArray...) = apply_element_wise(ps, params, input...)
     __pbs_executable(input::AbstractArray, params::NeuralNetworkParameters) = _pbs_executable(ps, params, input)
     # return this one if sinput & soutput are supplied
@@ -110,13 +110,13 @@ b = c(input, ps).^2
 (true, true)
 ```
 """
-function function_valued_parameters(eqs::NeuralNetworkParameters, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...)
-    vals = Tuple(build_nn_function(eqs[key], sparams, sinput...) for key in keys(eqs))
+function function_valued_parameters(eqs::NeuralNetworkParameters, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat)
+    vals = Tuple(build_nn_function(eqs[key], sparams, sinput...; reduce = reduce) for key in keys(eqs))
     NeuralNetworkParameters{keys(eqs)}(vals)
 end
 
-function function_valued_parameters(eqs::NamedTuple, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...)
-    vals = Tuple(build_nn_function(eqs[key], sparams, sinput...) for key in keys(eqs))
+function function_valued_parameters(eqs::NamedTuple, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat)
+    vals = Tuple(build_nn_function(eqs[key], sparams, sinput...; reduce = reduce) for key in keys(eqs))
     NamedTuple{keys(eqs)}(vals)
 end
 

src/utils/create_array.jl

@@ -0,0 +1,3 @@
+function Symbolics.SymbolicUtils.Code.create_array(::Type{<:Base.ReshapedArray{T, N, P}}, S, nd::Val, d::Val, elems...) where {T, N, P}
+    Symbolics.SymbolicUtils.Code.create_array(P, S, nd, d, elems...)
+end

test/reshape_test.jl

@@ -0,0 +1,40 @@
+using SymbolicNeuralNetworks
+using AbstractNeuralNetworks
+using Symbolics
+using Test
+
+function set_up_network()
+    c = Chain(Dense(2, 3))
+    nn = SymbolicNeuralNetwork(c)
+    soutput = nn.model(nn.input, nn.params)
+    nn_cpu = NeuralNetwork(c)
+    nn, soutput, nn_cpu
+end
+
+function test_for_input()
+    nn, soutput, nn_cpu = set_up_network()
+    input = rand(2, 5)
+    input2 = reshape((@view input[:, 1:2]), 2, 1, 2)
+    built_function = build_nn_function(soutput, nn.params, nn.input)
+    outputs = built_function(input2, nn_cpu.params)
+    for i in 1:2
+        @test nn.model(input[:, i], nn_cpu.params) ≈ outputs[:, 1, i]
+    end
+end
+
+function test_for_input_and_output()
+    nn, soutput2, nn_cpu = set_up_network()
+    input = rand(2, 5)
+    output = rand(3, 5)
+    input2 = reshape((@view input[:, 1:2]), 2, 1, 2)
+    output2 = reshape((@view output[:, 1:2]), 3, 1, 2)
+    @variables soutput[1:3]
+    built_function = build_nn_function((soutput - soutput2).^2, nn.params, nn.input, soutput)
+    outputs = built_function(input2, output2, nn_cpu.params)
+    for i in 1:2
+        @test (nn.model(input[:, i], nn_cpu.params) - output[:, i]).^2 ≈ outputs[:, 1, i]
+    end
+end
+
+test_for_input()
+test_for_input_and_output()

test/runtests.jl

@@ -3,7 +3,8 @@ using SafeTestsets
 using Test
 
 @safetestset "Docstring tests.                                                                      " begin include("doctest.jl") end
-@safetestset "Symbolic gradient                                                                      " begin include("symbolic_gradient.jl") end
-@safetestset "Symbolic Neural network                                                                " begin include("neural_network_derivative.jl") end
-@safetestset "Symbolic Params                                                                        " begin include("test_params.jl") end
-# @safetestset "HNN Loss                                                                               " begin include("test_hnn_loss_pullback.jl") end
+@safetestset "Symbolic gradient                                                                     " begin include("symbolic_gradient.jl") end
+@safetestset "Symbolic Neural network                                                               " begin include("neural_network_derivative.jl") end
+@safetestset "Symbolic Params                                                                       " begin include("test_params.jl") end
+# @safetestset "HNN Loss                                                                               " begin include("test_hnn_loss_pullback.jl") end
+@safetestset "Check if reshape works in the correct way with the generated functions.               " begin include("reshape_test.jl") end