Added test for reshape and create_array routine.

Author: benedict-96

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/utils/build_function.jl

@@ -25,9 +25,16 @@ end
 function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr)
     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)
+    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

@@ -19,8 +19,16 @@ end
 function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr)
     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)
+    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)
+        reshape(output_not_reshaped, size(output_not_reshaped, 1), size(x, 2), size(x, 3))
+    end
     gen_fun_returned
 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 input[:, 1:3]), 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