Reduce has to be optional (depending on whether we compute derivatives of neural network parameters or other expressions.

Author: benedict-96

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,9 +22,9 @@ 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, 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

src/utils/build_function2.jl

@@ -16,9 +16,9 @@ function build_nn_function(eqs, nn::AbstractSymbolicNeuralNetwork, soutput)
     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), hcat, axes(input, 2))
+    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
@@ -27,11 +27,20 @@ function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Sy
     # 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)
-        reshape(output_not_reshaped, size(output_not_reshaped, 1), size(x, 2), size(x, 3))
+        # 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