From 6521a747aff22e046e5fc5605d7c88553db24f1a Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 13:14:08 +0100
Subject: [PATCH 01/23] Moved build_function functions to separate directory.

---
 src/{utils => build_function}/build_function.jl        | 0
 src/{utils => build_function}/build_function2.jl       | 0
 src/{utils => build_function}/build_function_arrays.jl | 0
 3 files changed, 0 insertions(+), 0 deletions(-)
 rename src/{utils => build_function}/build_function.jl (100%)
 rename src/{utils => build_function}/build_function2.jl (100%)
 rename src/{utils => build_function}/build_function_arrays.jl (100%)

diff --git a/src/utils/build_function.jl b/src/build_function/build_function.jl
similarity index 100%
rename from src/utils/build_function.jl
rename to src/build_function/build_function.jl
diff --git a/src/utils/build_function2.jl b/src/build_function/build_function2.jl
similarity index 100%
rename from src/utils/build_function2.jl
rename to src/build_function/build_function2.jl
diff --git a/src/utils/build_function_arrays.jl b/src/build_function/build_function_arrays.jl
similarity index 100%
rename from src/utils/build_function_arrays.jl
rename to src/build_function/build_function_arrays.jl

From 6f59167de2b6d27c9d00bb3a8c8d3d102a56a0a6 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 13:15:22 +0100
Subject: [PATCH 02/23] Moved pullback.jl into derivatives directory.

---
 src/SymbolicNeuralNetworks.jl     | 8 ++++----
 src/{ => derivatives}/pullback.jl | 0
 2 files changed, 4 insertions(+), 4 deletions(-)
 rename src/{ => derivatives}/pullback.jl (100%)

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index c7d8294..de2c6a1 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -39,12 +39,12 @@ module SymbolicNeuralNetworks
     include("hamiltonian.jl")
 
     export build_nn_function
-    include("utils/build_function.jl")
-    include("utils/build_function2.jl")
-    include("utils/build_function_arrays.jl")
+    include("build_function/build_function.jl")
+    include("build_function/build_function2.jl")
+    include("build_function/build_function_arrays.jl")
 
     export SymbolicPullback
-    include("pullback.jl")
+    include("derivatives/pullback.jl")
 
     include("derivatives/derivative.jl")
     include("derivatives/jacobian.jl")
diff --git a/src/pullback.jl b/src/derivatives/pullback.jl
similarity index 100%
rename from src/pullback.jl
rename to src/derivatives/pullback.jl

From c659bbe85a6666e5b64b984966478f471fb05783 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 13:19:16 +0100
Subject: [PATCH 03/23] Moved files to new directory symbolic_neuralnet.

---
 src/SymbolicNeuralNetworks.jl                      | 4 ++--
 src/{ => symbolic_neuralnet}/symbolic_neuralnet.jl | 0
 src/{utils => symbolic_neuralnet}/symbolize.jl     | 0
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename src/{ => symbolic_neuralnet}/symbolic_neuralnet.jl (100%)
 rename src/{utils => symbolic_neuralnet}/symbolize.jl (100%)

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index de2c6a1..338123c 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -18,7 +18,7 @@ module SymbolicNeuralNetworks
     include("equation_types.jl")
 
     export symbolize
-    include("utils/symbolize.jl")
+    include("symbolic_neuralnet/symbolize.jl")
 
     export AbstractSymbolicNeuralNetwork
     export SymbolicNeuralNetwork, SymbolicModel
@@ -33,7 +33,7 @@ module SymbolicNeuralNetworks
     include("chain.jl")
 
     export evaluate_equations
-    include("symbolic_neuralnet.jl")
+    include("symbolic_neuralnet/symbolic_neuralnet.jl")
 
     export symbolic_hamiltonian
     include("hamiltonian.jl")
diff --git a/src/symbolic_neuralnet.jl b/src/symbolic_neuralnet/symbolic_neuralnet.jl
similarity index 100%
rename from src/symbolic_neuralnet.jl
rename to src/symbolic_neuralnet/symbolic_neuralnet.jl
diff --git a/src/utils/symbolize.jl b/src/symbolic_neuralnet/symbolize.jl
similarity index 100%
rename from src/utils/symbolize.jl
rename to src/symbolic_neuralnet/symbolize.jl

From c64a7d65b093e22923f1334a44fd82ba24658076 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 13:19:45 +0100
Subject: [PATCH 04/23] Made spacing uniform.

---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 8daf81c..af4c8f7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,7 +2,7 @@ using SymbolicNeuralNetworks
 using SafeTestsets
 using Test
 
-@safetestset "Docstring tests.                                                                      " begin include("doctest.jl") end
+@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

From 8acf5cb769ef1d0e57db45ee889fb25c9a585ce2 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 13:24:28 +0100
Subject: [PATCH 05/23] renamed utils to custom_definitions_and_extensions and
 moved relevant files to that directory. Also deleted files that are not
 needed anymore.

---
 src/SymbolicNeuralNetworks.jl                 |  9 +-
 .../equation_types.jl                         |  0
 .../latexraw.jl                               |  0
 src/custom_equation.jl                        | 89 -----------------
 src/hamiltonian.jl                            | 96 -------------------
 5 files changed, 2 insertions(+), 192 deletions(-)
 rename src/{ => custom_definitions_and_extensions}/equation_types.jl (100%)
 rename src/{utils => custom_definitions_and_extensions}/latexraw.jl (100%)
 delete mode 100644 src/custom_equation.jl
 delete mode 100644 src/hamiltonian.jl

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index 338123c..8183b9a 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -15,7 +15,7 @@ module SymbolicNeuralNetworks
     
     RuntimeGeneratedFunctions.init(@__MODULE__)
 
-    include("equation_types.jl")
+    include("custom_definitions_and_extensions/equation_types.jl")
 
     export symbolize
     include("symbolic_neuralnet/symbolize.jl")
@@ -35,9 +35,6 @@ module SymbolicNeuralNetworks
     export evaluate_equations
     include("symbolic_neuralnet/symbolic_neuralnet.jl")
 
-    export symbolic_hamiltonian
-    include("hamiltonian.jl")
-
     export build_nn_function
     include("build_function/build_function.jl")
     include("build_function/build_function2.jl")
@@ -50,7 +47,5 @@ module SymbolicNeuralNetworks
     include("derivatives/jacobian.jl")
     include("derivatives/gradient.jl")
 
-    include("custom_equation.jl")
-
-    include("utils/latexraw.jl")
+    include("custom_definitions_and_extensions/latexraw.jl")
 end
diff --git a/src/equation_types.jl b/src/custom_definitions_and_extensions/equation_types.jl
similarity index 100%
rename from src/equation_types.jl
rename to src/custom_definitions_and_extensions/equation_types.jl
diff --git a/src/utils/latexraw.jl b/src/custom_definitions_and_extensions/latexraw.jl
similarity index 100%
rename from src/utils/latexraw.jl
rename to src/custom_definitions_and_extensions/latexraw.jl
diff --git a/src/custom_equation.jl b/src/custom_equation.jl
deleted file mode 100644
index 542c9cd..0000000
--- a/src/custom_equation.jl
+++ /dev/null
@@ -1,89 +0,0 @@
-"""
-    substitute_gradient(eq, s∇nn, s∇output)
-
-Substitute the symbolic expression `s∇nn` in `eq` with the symbolic expression `s∇output`.
-
-# Implementation 
-
-See the comment in [`evaluate_equation`](@ref).
-"""
-function substitute_gradient(eq, s∇nn, s∇output)
-    @assert axes(s∇nn) == axes(s∇output)
-    substitute.(eq, Ref(Dict([s∇nn[i] => s∇output[i] for i in axes(s∇nn, 1)])))
-end
-
-function substitute_gradient(eq, ::Nothing, ::Nothing)
-    eq 
-end
-
-"""
-    evaluate_equation(eq, soutput)
-
-Replace `snn` in `eq` with `soutput` (input), scalarize and expand derivatives.
-
-# Implementation
-
-Here we use `Symbolics.substitute` with broadcasting to be able to handle `eq`s that are arrays.
-For that reason we use [`Ref` before `Dict`](https://discourse.julialang.org/t/symbolics-and-substitution-using-broadcasting/68705).
-This is also the case for the functions [`substitute_gradient`](@ref).
-"""
-function evaluate_equation(eq::EqT, snn::EqT, s∇nn::EqT, soutput::EqT, s∇output::EqT)
-    @assert axes(snn) == axes(soutput)
-    eq_output_substituted = substitute.(eq, Ref(Dict([snn[i] => soutput[i] for i in axes(snn, 1)])))
-    substitute_gradient(eq_output_substituted, s∇nn, s∇output)
-end
-
-"""
-    evaluate_equations(eqs, soutput)
-
-Apply [`evaluate_equation`](@ref) to a `NamedTuple` and append `(soutput = soutput, s∇output = s∇output)`.
-"""
-function evaluate_equations(eqs::NamedTuple, snn::EqT, s∇nn::EqT, soutput::EqT, s∇output::EqT; simplify = true)
-    
-    # closure
-    _evaluate_equation(eq) = evaluate_equation(eq, snn, s∇nn, soutput, s∇output)
-    evaluated_equations = Tuple(_evaluate_equation(eq) for eq in eqs)
-
-    soutput_eq = (soutput = simplify == true ? Symbolics.simplify(soutput) : soutput,)
-    s∇output_eq = isnothing(s∇nn) ? NamedTuple() : (s∇output = simplify == true ? Symbolics.simplify(s∇output) : s∇output,)
-    merge(soutput_eq, s∇output_eq, NamedTuple{keys(eqs)}(evaluated_equations))
-end
-
-"""
-    evaluate_equations(eqs, nn)
-
-Expand the output and gradient in `eqs` with the weights in `nn`.
-
-- `eqs` here has to be a `NamedTuple` that contains keys 
-- `:x`: gives the inputs to the neural network and 
-- `:nn`: symbolic expression of the neural network.
-
-# Implementation
-
-Internally this 
-1. computes the gradient and
-2. calls [`evaluate_equations(::NamedTuple, ::EqT, ::EqT, ::EqT, EqT)`](@ref).
-
-"""
-function evaluate_equations(eqs::NamedTuple, nn::AbstractSymbolicNeuralNetwork; kwargs...)
-    @assert [:x, :nn] ⊆ keys(eqs)
-    
-    sinput = eqs.x
-
-    snn = eqs.nn
-
-    s∇nn = haskey(eqs, :∇nn) ? eqs.∇nn : nothing
-
-    remaining_eqs = NamedTuple([p for p in pairs(eqs) if p[1] ∉ [:x, :nn, :∇nn]])
-
-    # Evaluation of the symbolic output
-    soutput = _scalarize(nn.model(sinput, nn.params))
-
-    # make differential 
-    Dx = symbolic_differentials(sinput)
-
-    # Evaluation of gradient
-    s∇output = isnothing(s∇nn) ? nothing : [expand_derivatives(Symbolics.scalarize(dx(soutput))) for dx in Dx]
-
-    evaluate_equations(remaining_eqs, snn, s∇nn, soutput, s∇output; kwargs...)
-end
\ No newline at end of file
diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl
deleted file mode 100644
index 722ba74..0000000
--- a/src/hamiltonian.jl
+++ /dev/null
@@ -1,96 +0,0 @@
-"""
-    HamiltonianSymbolicNeuralNetwork <: AbstractSymbolicNeuralNetwork
-
-A struct that inherits properties from the abstract type `AbstractSymbolicNeuralNetwork`.
-
-# Constructor
-
-    HamiltonianSymbolicNeuralNetwork(model)
-
-Make an instance of `HamiltonianSymbolicNeuralNetwork` based on a `Chain` or an `Architecture`.
-This is similar to the constructor for [`SymbolicNeuralNetwork`](@ref) but also checks if the input dimension is even-dimensional and the output dimension is one.
-"""
-struct HamiltonianSymbolicNeuralNetwork{AT, MT, PT} <: AbstractSymbolicNeuralNetwork{AT}
-    architecture::AT
-    model::MT
-    params::PT
-end
-
-function HamiltonianSymbolicNeuralNetwork(arch::Architecture, model::Model)
-    @assert iseven(input_dimension(model)) "Input dimension has to be an even number."
-    @assert output_dimension(model) == 1 "Output dimension of network has to be scalar."
-
-    sparams = symbolicparameters(model)
-    HamiltonianSymbolicNeuralNetwork(arch, model, sparams)
-end
-
-HamiltonianSymbolicNeuralNetwork(model::Model) = HamiltonianSymbolicNeuralNetwork(UnknownArchitecture(), model)
-HamiltonianSymbolicNeuralNetwork(arch::Architecture) = HamiltonianSymbolicNeuralNetwork(arch, Chain(model))
-
-"""
-    vector_field(nn::HamiltonianSymbolicNeuralNetwork)
-
-Get the symbolic expression for the vector field belonging to the HNN `nn`.
-
-# Implementation 
-
-This is calling [`SymbolicNeuralNetworks.Jacobian`](@ref) and then multiplies the result with a Poisson tensor.
-"""
-function vector_field(nn::HamiltonianSymbolicNeuralNetwork)
-    gradient_output = gradient(nn)
-    sinput, soutput, ∇nn = gradient_output.x, gradient_output.soutput, gradient_output.s∇output
-    input_dim = input_dimension(nn.model)
-    n = input_dim ÷ 2
-    # placeholder for one
-    @variables o
-    o_vec = repeat([o], n)
-    𝕀 = Diagonal(o_vec)
-    𝕆 = zero(𝕀)
-    𝕁 = hcat(vcat(𝕆, -𝕀), vcat(𝕀, 𝕆))
-    (x = sinput, nn = soutput, ∇nn = ∇nn, hvf = substitute(𝕁 * ∇nn, Dict(o => 1, )))
-end
-
-"""
-    HNNLoss <: NetworkLoss
-
-The loss for a Hamiltonian neural network.
-
-# Constructor
-
-This can be called with an instance of [`HamiltonianSymbolicNeuralNetwork`](@ref) as the only input arguemtn, i.e.:
-```julia
-HNNLoss(nn)
-```
-where `nn` is a [`HamiltonianSymbolicNeuralNetwork`](@ref) gives the corresponding Hamiltonian loss.
-
-# Funktor
-
-```julia
-loss(c, ps, input, output)
-loss(ps, input, output) # equivalent to the above
-```
-"""
-struct HNNLoss{FT} <: NetworkLoss
-    hvf::FT
-end
-
-function HNNLoss(nn::HamiltonianSymbolicNeuralNetwork)
-    x_hvf = vector_field(nn)
-    x = x_hvf.x
-    hvf = x_hvf.hvf
-    hvf_function = build_nn_function(hvf, x, nn)
-    HNNLoss(hvf_function)
-end
-
-function (loss::HNNLoss)(   ::Union{Chain, AbstractExplicitLayer}, 
-                            ps::Union{NeuralNetworkParameters, NamedTuple}, 
-                            input::QPTOAT, 
-                            output::QPTOAT)
-    loss(ps, input, output)
-end
-
-function (loss::HNNLoss)(   ps::Union{NeuralNetworkParameters, NamedTuple},
-                            input::QPTOAT,
-                            output::QPTOAT)
-    norm(loss.hvf(input, ps) - output) / norm(output)
-end
\ No newline at end of file

From 7789dac99e5e22ea924cde10b1afe747d1f941dc Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 13:26:13 +0100
Subject: [PATCH 06/23] Removed HamiltonianNN use.

---
 src/derivatives/pullback.jl | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/derivatives/pullback.jl b/src/derivatives/pullback.jl
index 31da2b3..dba9c53 100644
--- a/src/derivatives/pullback.jl
+++ b/src/derivatives/pullback.jl
@@ -85,10 +85,6 @@ struct SymbolicPullback{NNLT, FT} <: AbstractPullback{NNLT}
     fun::FT
 end
 
-function SymbolicPullback(nn::HamiltonianSymbolicNeuralNetwork)
-    SymbolicPullback(nn, HNNLoss(nn))
-end
-
 function SymbolicPullback(nn::SymbolicNeuralNetwork, loss::NetworkLoss)
     @variables soutput[1:output_dimension(nn.model)]
     symbolic_loss = loss(nn.model, nn.params, nn.input, soutput)

From 7b332c22e4c9e884b595731fc4f8f321d71d9147 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 14:06:15 +0100
Subject: [PATCH 07/23] Moved chain.jl in separate directory.

---
 src/SymbolicNeuralNetworks.jl | 2 +-
 src/{ => chain}/chain.jl      | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename src/{ => chain}/chain.jl (100%)

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index 8183b9a..c7b8c07 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -30,7 +30,7 @@ module SymbolicNeuralNetworks
     include("layers/abstract.jl")
     include("layers/dense.jl")
     include("layers/linear.jl")
-    include("chain.jl")
+    include("chain/chain.jl")
 
     export evaluate_equations
     include("symbolic_neuralnet/symbolic_neuralnet.jl")
diff --git a/src/chain.jl b/src/chain/chain.jl
similarity index 100%
rename from src/chain.jl
rename to src/chain/chain.jl

From e438730b8fdbbb716d1490246398def0ff7ecbb3 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 14:06:29 +0100
Subject: [PATCH 08/23] Removed tests that aren't really needed anymore.

---
 test/comparison_performace.jl  |  61 -------------
 test/plots.jl                  | 152 ---------------------------------
 test/runtests.jl               |   3 +-
 test/test_allocations.jl       | 100 ----------------------
 test/test_de_envelop.jl        |  15 ----
 test/test_debug_sympnet.jl     |  80 -----------------
 test/test_get_track.jl         |   2 -
 test/test_hnn_loss_pullback.jl |  33 -------
 test/test_rewrite.jl           |  19 -----
 9 files changed, 1 insertion(+), 464 deletions(-)
 delete mode 100644 test/comparison_performace.jl
 delete mode 100644 test/plots.jl
 delete mode 100644 test/test_allocations.jl
 delete mode 100644 test/test_de_envelop.jl
 delete mode 100644 test/test_debug_sympnet.jl
 delete mode 100644 test/test_get_track.jl
 delete mode 100644 test/test_hnn_loss_pullback.jl
 delete mode 100644 test/test_rewrite.jl

diff --git a/test/comparison_performace.jl b/test/comparison_performace.jl
deleted file mode 100644
index 6e735c6..0000000
--- a/test/comparison_performace.jl
+++ /dev/null
@@ -1,61 +0,0 @@
-using GeometricMachineLearning
-using SymbolicNeuralNetworks
-using Symbolics
-using GeometricEquations
-using Test
-include("plots.jl")
-
-using GeometricProblems.HarmonicOscillator
-using GeometricProblems.HarmonicOscillator: hodeensemble, hamiltonian, default_parameters
-
-
-# Importing Data
-ensemble_problem = hodeensemble(tspan = (0.0,4.0))
-ensemble_solution =  exact_solution(ensemble_problem ) 
-training_data = TrainingData(ensemble_solution)
-
-# Extension of symbolicparameters method for Gradient Layer
-
-function SymbolicNeuralNetworks.symbolicparameters(d::Gradient{M, N, true}) where {M,N}
-    @variables K[1:d.second_dim÷2, 1:M÷2]
-    @variables b[1:d.second_dim÷2]
-    @variables a[1:d.second_dim÷2]
-    (weight = K, bias = b, scale = a)
-end
-
-function SymbolicNeuralNetworks.symbolicparameters(d::Gradient{M, N, false}) where {M,N}
-    @variables a[1:d.second_dim÷2, 1:1]
-    (scale = a, )
-end
-
-#creating of the neuralnetwork and the symbolized one
-arch = GSympNet(2; nhidden = 4, width = 10, allow_fast_activation = false)
-sympnet = NeuralNetwork(arch, Float64)
-ssympnet = symbolize(sympnet, 2)
-
-#parameters for the training
-method = BasicSympNet()
-mopt = AdamOptimizer()
-
-
-function performance_symbolic(nruns::Int)
-    training_parameters =TrainingParameters(nruns, method, mopt)
-    training_set = TrainingSet(ssympnet, training_parameters, training_data)
-    train!(training_set; showprogress = true, timer = true)
-    nothing
-end
-
-function performance_withoutsymbolic(nruns::Int)
-    training_parameters =TrainingParameters(nruns, method, mopt)
-    training_set = TrainingSet(sympnet, training_parameters, training_data)
-    train!(training_set; showprogress = true, timer = true)
-    nothing
-end
-
-
-#Plots
-
-function _plot(neuralnetsolution)
-    H(x) = hamiltonian(x[1+length(x)÷2:end], 0.0, x[1:length(x)], default_parameters)
-    plot_result(training_data, neural_net_solution, H; batch_nb_trajectory = 10, filename = "GSympNet 4-10 on Harmonic Oscillator", nb_prediction = 5)
-end
\ No newline at end of file
diff --git a/test/plots.jl b/test/plots.jl
deleted file mode 100644
index 90b3863..0000000
--- a/test/plots.jl
+++ /dev/null
@@ -1,152 +0,0 @@
-using Plots
-using LaTeXStrings
-using StatsBase
-
-
-
-
-function plot_data(data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, title::String = ""; index::AbstractArray = 1:get_nb_trajectory(data))
-
-    plt = plot(size=(1000,1000), titlefontsize=15, guidefontsize=14)
-
-    for i in index
-        plot!(vcat([get_data(data,:q, i, n) for n in 1:get_length_trajectory(data,i)]...), vcat([get_data(data,:p, i,n) for n in 1:get_length_trajectory(data,i)]...), label="Training data "*string(i),linewidth = 3,mk=*)
-    end
-
-    title!(title)
-
-    xlabel!(L"q")
-    xlims!((-3.5,3.5))
-    ylabel!(L"p")
-    ylims!((-2.5,2.5))
-
-    plot!(legend=:outerbottom,legendcolumns=2)
-
-    return plt
-     
-end
-
-
-function plot_verification(data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, nns::NeuralNetSolution; index::AbstractArray = [1])
-
-    plt = plot(size=(1000,800), titlefontsize=15, legendfontsize=10, guidefontsize=14)
-
-    for i in index
-        plot!(vcat([get_data(data,:q, i, n) for n in 1:get_length_trajectory(data,i)]...), vcat([get_data(data,:p, i,n) for n in 1:get_length_trajectory(data,i)]...), label="Training data "*string(i),linewidth = 3,mk=*)
-        q = []
-        p = []
-        qp = [get_data(data,:q,i,1)..., get_data(data,:p,i,1)...]
-        push!(q,qp[1])
-        push!(p,qp[2])
-        for _ in 2:get_length_trajectory(data,i)
-            qp = nns.nn(qp)
-            push!(q,qp[1])
-            push!(p,qp[2])
-        end
-        scatter!(q,p, label="Learned trajectory "*string(i), mode="markers+lines", ma = 0.8)
-    end
-
-    xlabel!(L"q")
-    xlims!((-3.5,3.5))
-    ylabel!(L"p")
-    ylims!((-2.5,2.5))
-
-    title!("Verifications")
-
-    plot!(legend=:outerbottom,legendcolumns=2)
-
-    return plt
-end
-
-
-function plot_loss()
-
-    
-
-end
-
-
-
-
-function plot_prediction(data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, nns::NeuralNetSolution, initial_cond::AbstractArray, H; scale = 1)
-
-    plt = plot(size=(1000,800), titlefontsize=15, legendfontsize=10, guidefontsize=14)
-
-    xmin = -3.5*scale
-    xmax = 3.5*scale
-    ymin = -2.5*scale
-    ymax = 2.5*scale
-    xlabel!(L"q")
-    xlims!((xmin,xmax))
-    ylabel!(L"p")
-    ylims!((ymin,ymax))
-
-    X = range(xmin, stop=xmax, length=100)
-    Y = range(ymin, stop=ymax, length=100)
-    contour!(X, Y, [H([x,y]) for y in Y, x in X], linewidth = 0, fill = true, levels = 7, c = cgrad(:default, rev = true))
-
-
-
-    arrow_indices = [10, 20, 30, 40, 50, 60, 70, 80, 90]
-    i=0
-    for qp0 in initial_cond
-        i+=1
-        q = []
-        p = []
-        qp = qp0
-        push!(q,qp[1])
-        push!(p,qp[2])
-        for _ in 2:100
-            qp = nns.nn(qp)
-            push!(q,qp[1])
-            push!(p,qp[2])
-        end
-        scatter!(q,p, label="Prediction "*string(i), mode="markers+lines", ma = 0.8)
-        #quiver!(q[arrow_indices], p[arrow_indices], quiver=(0.2, 0.2, :auto))
-    end
-
-
-
-    title!("Predictions")
-
-    plot!(legend=:outerbottom,legendcolumns=2)
-
-    return plt
-end
-
-function plot_result(data::TrainingData, nns::NeuralNetSolution, hamiltonian; batch_nb_trajectory::Int = get_nb_trajectory(data), batch_verif::Int = 3, filename = nothing, nb_prediction = 2)
-
-    plt_data = plot_data(data, "Datas"; index = sort!(sample(1:get_nb_trajectory(data), batch_nb_trajectory, replace = false)))
-
-    plt_verif = plot_verification(data, nns; index = sort!(sample(1:get_nb_trajectory(data), batch_verif, replace = false)))
-
-    initial_conditions = [(q = get_data(data,:q,i,1), p = get_data(data,:p,i,1)) for i in 1:get_nb_trajectory(data)]
-    min_q = min([initial_conditions[i][:q] for i in 1:get_nb_trajectory(data)]...)
-    min_p = min([initial_conditions[i][:p] for i in 1:get_nb_trajectory(data)]...)
-    max_q = max([initial_conditions[i][:q] for i in 1:get_nb_trajectory(data)]...)
-    max_p = max([initial_conditions[i][:p] for i in 1:get_nb_trajectory(data)]...)
-
-    initial_cond = [[linear_trans(rand(), min_q, max_q)..., linear_trans(rand(), min_p, max_p)...] for _ in 1:nb_prediction]
-
-    plt_pred = plot_prediction(data, nns, initial_cond, hamiltonian)
-
-    initial_cond_far = [[linear_trans(rand(), 10*min_q, 10*max_q)..., linear_trans(rand(), 10*min_p, 10*max_p)...] for _ in 1:nb_prediction]
-
-    plt_farpred = plot_prediction(data, nns, initial_cond_far, hamiltonian; scale = 10)
-
-
-    l = @layout grid(2, 2)
-
-    plt = plot(plt_data, plt_verif, plt_pred, plt_farpred, layout = l)
-
-
-    if filename !== nothing
-        savefig(filename)
-    end
-
-    return plt
-end
-
-
-
-linear_trans(x,a,b) = x * (b-a) + a
diff --git a/test/runtests.jl b/test/runtests.jl
index af4c8f7..4891f9c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,5 +5,4 @@ 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
\ No newline at end of file
+@safetestset "Symbolic Params                                                                        " begin include("test_params.jl") end
\ No newline at end of file
diff --git a/test/test_allocations.jl b/test/test_allocations.jl
deleted file mode 100644
index 6e11a79..0000000
--- a/test/test_allocations.jl
+++ /dev/null
@@ -1,100 +0,0 @@
-using Symbolics
-using GeometricMachineLearning
-using SymbolicNeuralNetworks
-using RuntimeGeneratedFunctions
-#using BenchmarkTools
-
-RuntimeGeneratedFunctions.init(@__MODULE__)
-
-function SymbolicNeuralNetworks.symbolicparameters(d::Gradient{M, N, true}) where {M,N}
-    @variables K[1:d.second_dim÷2, 1:M÷2]
-    @variables b[1:d.second_dim÷2]
-    @variables a[1:d.second_dim÷2]
-    (weight = K, bias = b, scale = a)
-end
-
-function SymbolicNeuralNetworks.symbolicparameters(d::Gradient{M, N, false}) where {M,N}
-    @variables a[1:d.second_dim÷2, 1:1]
-    (scale = a, )
-end
-
-arch = GSympNet(2; nhidden = 1, width = 4, allow_fast_activation = false)
-sympnet = NeuralNetwork(arch, Float64)
-
-ssympnet = SymbolicNeuralNetwork(arch, 2)
-
-eva = equations(ssympnet).eval
-
-@variables x[1:2]
-sparams = symbolicparameters(model(ssympnet))
-
-code = build_function(eva, x, sparams...)[1]
-
-postcode = SymbolicNeuralNetworks.rewrite_neuralnetwork(code, (x,), sparams)
-
-x = [1,2]
-
-H = functions(ssympnet).eval
-
-expr = :(function (x::AbstractArray, params::Tuple)
-    SymbolicUtils.Code.create_array(Array, nothing, Val{1}(), Val{(2,)}(), getindex(broadcast(+, Real[x[1]], adjoint((params[1]).weight) * broadcast(*, (params[1]).scale, broadcast(tanh, broadcast(+, (params[1]).weight * Real[x[2]], (params[1]).bias)))), 1), getindex(x, 2))
-end)
-#=
-SymbolicUtils.Code.create_array(Array, nothing, Val{1}(), Val{(2,)}(), getindex(broadcast(+, Real[x[1]], adjoint((params[1]).weight) * broadcast(*, (params[1]).scale, broadcast(tanh, broadcast(+, (params[1]).weight * Real[x[2]], (params[1]).bias)))), 1), getindex(x, 2))
-=#
-
-expr2 = :(function (x::AbstractArray, p::Tuple)
-    getindex(broadcast(+, x[1], adjoint((p[1]).weight) * broadcast(*, (p[1]).scale, broadcast(tanh, broadcast(+, (p[1]).weight * x[2], (p[1]).bias)))), 1), getindex(x, 2)
-end)
-
-fun1 = @RuntimeGeneratedFunction(Symbolics.inject_registered_module_functions(expr))
-fun1(x, sympnet.params)
-
-fun2 = @RuntimeGeneratedFunction(Symbolics.inject_registered_module_functions(expr2))
-fun2(x, sympnet.params)
-
-sympnet(x, sympnet.params)
-
-@time functions(ssympnet).eval(x, sympnet.params)
-@time fun1(x, sympnet.params)
-@time fun2(x, sympnet.params)
-@time sympnet(x, sympnet.params)
-
-#=
-@benchmark functions(ssympnet).eval(x, sympnet.params)
-@benchmark fun1(x, sympnet.params)
-@benchmark fun2(x, sympnet.params)
-@benchmark sympnet(x, sympnet.params)
-=#
-
-function optimize_code!(expr)
-    try expr.args
-    catch
-        return expr
-    end
-    for i in eachindex(expr.args)
-        expr.args[i] =  optimize_code!(expr.args[i])
-    end
-    if expr.args[1] == :broadcast
-        if length(expr.args) == 4
-            return :(($(expr.args[2])).($(expr.args[3]), $(expr.args[4])))
-        elseif length(expr.args) == 3
-            return :(($(expr.args[2])).($(expr.args[3])))
-        end
-    elseif expr.args[1] == :getindex
-        return Meta.parse(string(expr.args[2],"[",expr.args[3],"]"))
-    elseif expr.args[1] == :Real
-        return expr.args[2]
-    end
-    return expr
-end
-
-expr = optimize_code!(expr)
-
-expr = Meta.parse(replace(string(expr), "SymbolicUtils.Code.create_array(typeof(sinput), nothing, Val{1}(), Val{(2,)}()," => "(" ))
-
-
-func = @RuntimeGeneratedFunction(Symbolics.inject_registered_module_functions(expr))
-func(x, sympnet.params)
-
-@time func(x, sympnet.params)
diff --git a/test/test_de_envelop.jl b/test/test_de_envelop.jl
deleted file mode 100644
index a615183..0000000
--- a/test/test_de_envelop.jl
+++ /dev/null
@@ -1,15 +0,0 @@
-using SymbolicNeuralNetworks
-using Test
-
-S = ((W=[1,2], b = 7), ((4,5), 7))
-dS = develop(S)
-Es = envelop(S, dS)[1]
-
-@test dS == [[1,2], 7, 4, 5, 7]
-@test S == Es
-
-dSc = develop(S; completely = true)
-Esc = envelop(S, dSc; completely = true)[1]
-
-@test dSc == [1,2, 7, 4, 5, 7]
-@test S == Esc
\ No newline at end of file
diff --git a/test/test_debug_sympnet.jl b/test/test_debug_sympnet.jl
deleted file mode 100644
index a14e390..0000000
--- a/test/test_debug_sympnet.jl
+++ /dev/null
@@ -1,80 +0,0 @@
-using GeometricMachineLearning
-using SymbolicNeuralNetworks
-using Symbolics
-using KernelAbstractions
-using Test
-
-function SymbolicNeuralNetworks.symbolicparameters(d::Gradient{M, N, true}) where {M,N}
-    @variables K[1:d.second_dim÷2, 1:M÷2]
-    @variables b[1:d.second_dim÷2]
-    @variables a[1:d.second_dim÷2]
-    (weight = K, bias = b, scale = a)
-end
-
-function SymbolicNeuralNetworks.symbolicparameters(d::Gradient{M, N, false}) where {M,N}
-    @variables a[1:d.second_dim÷2, 1:1]
-    (scale = a, )
-end
-
-arch = GSympNet(2; nhidden = 1, width = 4, allow_fast_activation = false)
-sympnet=NeuralNetwork(arch, Float64)
-
-ssympnet = SymbolicNeuralNetwork(arch, 2)
-
-x = [1,2]
-@test functions(ssympnet).eval(x, sympnet.params) == sympnet(x)
-
-
-@time functions(ssympnet).eval(x, sympnet.params)
-@time sympnet(x)
-
-
-
-
-
-#=
-@kernel function assign_first_half!(q::AbstractVector, x::AbstractVector)
-    i = @index(Global)
-    q[i] = x[i]
-end
-
-@kernel function assign_second_half!(p::AbstractVector, x::AbstractVector, N::Integer)
-    i = @index(Global)
-    p[i] = x[i+N]
-end
-
-@kernel function assign_first_half!(q::AbstractMatrix, x::AbstractMatrix)
-    i,j = @index(Global, NTuple)
-    q[i,j] = x[i,j]
-end
-
-@kernel function assign_second_half!(p::AbstractMatrix, x::AbstractMatrix, N::Integer)
-    i,j = @index(Global, NTuple)
-    p[i,j] = x[i+N,j]
-end
-
-@kernel function assign_first_half!(q::AbstractArray{T, 3}, x::AbstractArray{T, 3}) where T 
-    i,j,k = @index(Global, NTuple)
-    q[i,j,k] = x[i,j,k]
-end
-
-@kernel function assign_second_half!(p::AbstractArray{T, 3}, x::AbstractArray{T, 3}, N::Integer) where T 
-    i,j,k = @index(Global, NTuple)
-    p[i,j,k] = x[i+N,j,k]
-end
-
-function assign_q_and_p(x::AbstractVector, N)
-    backend = try KernelAbstractions.get_backend(x)
-    catch 
-        CPU() end
-    q = KernelAbstractions.allocate(backend, eltype(x), N)
-    p = KernelAbstractions.allocate(backend, eltype(x), N)
-    q_kernel! = assign_first_half!(backend)
-    p_kernel! = assign_second_half!(backend)
-    q_kernel!(q, x, ndrange=size(q))
-    p_kernel!(p, x, N, ndrange=size(p))
-    (q, p)
-end
-
-assign_q_and_p(x, 1)
-=#
\ No newline at end of file
diff --git a/test/test_get_track.jl b/test/test_get_track.jl
deleted file mode 100644
index 690711b..0000000
--- a/test/test_get_track.jl
+++ /dev/null
@@ -1,2 +0,0 @@
-using SymbolicNeuralNetworks
-using Test
\ No newline at end of file
diff --git a/test/test_hnn_loss_pullback.jl b/test/test_hnn_loss_pullback.jl
deleted file mode 100644
index ec54cbc..0000000
--- a/test/test_hnn_loss_pullback.jl
+++ /dev/null
@@ -1,33 +0,0 @@
-using SymbolicNeuralNetworks
-using AbstractNeuralNetworks
-using GeometricMachineLearning: ZygotePullback
-using Symbolics
-using Test
-using Zygote
-
-function test_hnn_loss(input_dim::Integer = 2, nhidden::Integer = 1, hidden_dim::Integer = 3, T::DataType = Float64, second_axis = 1, third_axis = 1)
-    c = Chain(Dense(input_dim, hidden_dim, tanh), Tuple(Dense(hidden_dim, hidden_dim, tanh) for _ in 1:nhidden)..., Dense(hidden_dim, 1, identity; use_bias = false))
-    nn = NeuralNetwork(c, T)
-    snn = HamiltonianSymbolicNeuralNetwork(c)
-    loss = HNNLoss(snn)
-    zpb = ZygotePullback(loss)
-    spb = SymbolicPullback(snn)
-    input = rand(T, input_dim, second_axis, third_axis)
-    output = rand(T, input_dim, second_axis, third_axis)
-    # the x -> x[1].params is necessary because of Zygote idiosyncrasies
-    zpb_evaluated = zpb(nn.params, c, (input, output))[2](1)[1].params
-    spb_evaluated = spb(nn.params, c, (input, output))[2](1)
-    @assert keys(zpb_evaluated) == keys(spb_evaluated)
-    for key in keys(zpb_evaluated) @assert keys(zpb_evaluated[key]) == keys(spb_evaluated[key]) end
-    for key1 in keys(zpb_evaluated) for key2 in keys(zpb_evaluated[key1]) @test zpb_evaluated[key1][key2] ≈ spb_evaluated[key1][key2] end end
-end
-
-for input_dim in (2, )
-    for nhidden in (1, )
-        for hidden_dim in (2, )
-            for T in (Float32, Float64)
-                test_hnn_loss(input_dim, nhidden, hidden_dim, T)
-            end
-        end
-    end
-end
\ No newline at end of file
diff --git a/test/test_rewrite.jl b/test/test_rewrite.jl
deleted file mode 100644
index 991be5e..0000000
--- a/test/test_rewrite.jl
+++ /dev/null
@@ -1,19 +0,0 @@
-using SymbolicNeuralNetworks
-using Symbolics
-using Test
-
-@variables W1[1:2,1:2] W2[1:2,1:2] b1[1:2] b2[1:2]
-sparams = ((W =  W1, b = b1), (W = W2, b = b2))
-
-@variables st
-@variables sargs(st)[1:2]
-
-output = sparams[2].W  * tanh.(sparams[1].W * sargs + sparams[1].b) + sparams[2].b
-
-code_output = build_function(output, sargs..., develop(sparams)...)[2]
-rewrite_output = eval(rewrite_code(code_output, Tuple(sargs), sparams, "OUTPUT"))
-
-params = ((W = [1 3; 2 2], b = [1, 0]), (W = [1 1; 0 2], b = [1, 0]))
-args = [1, 0.2]
-
-@test_nowarn rewrite_output(args, params) 

From 1b5f9e45ccaa31d80301b64e1f693a28e60dbc49 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 14:20:55 +0100
Subject: [PATCH 09/23] Adjusted tests layout to how they appear in src.

---
 .../jacobian.jl}                                            | 0
 test/{ => derivatives}/symbolic_gradient.jl                 | 0
 test/runtests.jl                                            | 6 +++---
 test/{test_params.jl => symbolic_neuralnet/symbolize.jl}    | 0
 4 files changed, 3 insertions(+), 3 deletions(-)
 rename test/{neural_network_derivative.jl => derivatives/jacobian.jl} (100%)
 rename test/{ => derivatives}/symbolic_gradient.jl (100%)
 rename test/{test_params.jl => symbolic_neuralnet/symbolize.jl} (100%)

diff --git a/test/neural_network_derivative.jl b/test/derivatives/jacobian.jl
similarity index 100%
rename from test/neural_network_derivative.jl
rename to test/derivatives/jacobian.jl
diff --git a/test/symbolic_gradient.jl b/test/derivatives/symbolic_gradient.jl
similarity index 100%
rename from test/symbolic_gradient.jl
rename to test/derivatives/symbolic_gradient.jl
diff --git a/test/runtests.jl b/test/runtests.jl
index 4891f9c..c094b2b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -3,6 +3,6 @@ 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
\ No newline at end of file
+@safetestset "Symbolic gradient                                                                      " begin include("derivatives/symbolic_gradient.jl") end
+@safetestset "Symbolic Neural network                                                                " begin include("derivatives/jacobian.jl") end
+@safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end
\ No newline at end of file
diff --git a/test/test_params.jl b/test/symbolic_neuralnet/symbolize.jl
similarity index 100%
rename from test/test_params.jl
rename to test/symbolic_neuralnet/symbolize.jl

From 3ff82f4bcb8917402dff90221fd36025af9a96ad Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Fri, 13 Dec 2024 16:54:08 +0100
Subject: [PATCH 10/23] 1 -> begin.

---
 src/symbolic_neuralnet/symbolic_neuralnet.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/symbolic_neuralnet/symbolic_neuralnet.jl b/src/symbolic_neuralnet/symbolic_neuralnet.jl
index ae53736..73ec9be 100644
--- a/src/symbolic_neuralnet/symbolic_neuralnet.jl
+++ b/src/symbolic_neuralnet/symbolic_neuralnet.jl
@@ -50,7 +50,7 @@ end
 apply(snn::AbstractSymbolicNeuralNetwork, x, args...) = snn(x, args...)
 
 input_dimension(::AbstractExplicitLayer{M}) where M = M 
-input_dimension(c::Chain) = input_dimension(c.layers[1])
+input_dimension(c::Chain) = input_dimension(c.layers[begin])
 output_dimension(::AbstractExplicitLayer{M, N}) where {M, N} = N
 output_dimension(c::Chain) = output_dimension(c.layers[end])
 

From 7683404b7b2333c83f6bb5808d9d6b690b726a39 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Sat, 14 Dec 2024 20:41:38 +0100
Subject: [PATCH 11/23] Introduced a new constructor for SymbolicNeuralNetwork
 (using Neural Network from AbstractNeuralNetworks) and simplified symbolic
 parameter generation to a single function 'symbolize\!'.

---
 src/SymbolicNeuralNetworks.jl                |   7 +-
 src/build_function/build_function.jl         |  15 +-
 src/derivatives/gradient.jl                  |   4 +-
 src/layers/abstract.jl                       |   6 +-
 src/symbolic_neuralnet/symbolic_neuralnet.jl |  15 +-
 src/symbolic_neuralnet/symbolize.jl          | 154 ++++++++++++++-----
 test/derivatives/symbolic_gradient.jl        |  28 ++--
 test/symbolic_neuralnet/symbolize.jl         |  29 ++--
 8 files changed, 173 insertions(+), 85 deletions(-)

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index c7b8c07..b07d4c0 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -17,22 +17,17 @@ module SymbolicNeuralNetworks
 
     include("custom_definitions_and_extensions/equation_types.jl")
 
-    export symbolize
     include("symbolic_neuralnet/symbolize.jl")
 
     export AbstractSymbolicNeuralNetwork
-    export SymbolicNeuralNetwork, SymbolicModel
-    export HamiltonianSymbolicNeuralNetwork, HNNLoss
-    export architecture, model, params, equations, functions
+    export SymbolicNeuralNetwork
 
     # make symbolic parameters (`NeuralNetworkParameters`)
-    export symbolicparameters
     include("layers/abstract.jl")
     include("layers/dense.jl")
     include("layers/linear.jl")
     include("chain/chain.jl")
 
-    export evaluate_equations
     include("symbolic_neuralnet/symbolic_neuralnet.jl")
 
     export build_nn_function
diff --git a/src/build_function/build_function.jl b/src/build_function/build_function.jl
index 33d5469..33d532f 100644
--- a/src/build_function/build_function.jl
+++ b/src/build_function/build_function.jl
@@ -39,19 +39,18 @@ Build a function that can process a matrix. This is used as a starting point for
 # Examples
 
 ```jldoctest
-using SymbolicNeuralNetworks: _build_nn_function, symbolicparameters
-using Symbolics
+using SymbolicNeuralNetworks: _build_nn_function
+using SymbolicNeuralNetworks
 using AbstractNeuralNetworks
 
 c = Chain(Dense(2, 1, tanh))
-params = symbolicparameters(c)
-@variables sinput[1:2]
-eq = c(sinput, params)
-built_function = _build_nn_function(eq, params, sinput)
-ps = initialparameters(c)
+nn = NeuralNetwork(c)
+snn = SymbolicNeuralNetwork(nn)
+eq = c(snn.input, snn.params)
+built_function = _build_nn_function(eq, snn.params, snn.input)
 input = rand(2, 2)
 
-(built_function(input, ps, 1), built_function(input, ps, 2)) .≈ (c(input[:, 1], ps), c(input[:, 2], ps))
+(built_function(input, nn.params, 1), built_function(input, nn.params, 2)) .≈ (c(input[:, 1], nn.params), c(input[:, 2], nn.params))
 
 # output
 
diff --git a/src/derivatives/gradient.jl b/src/derivatives/gradient.jl
index b075201..ea519ec 100644
--- a/src/derivatives/gradient.jl
+++ b/src/derivatives/gradient.jl
@@ -29,7 +29,7 @@ nn = SymbolicNeuralNetwork(c)
 L"\begin{equation}
 \left[
 \begin{array}{c}
-1 - \tanh^{2}\left( \mathtt{b\_1}_{1} + \mathtt{W\_1}_{1,1} \mathtt{sinput}_{1} + \mathtt{W\_1}_{1,2} \mathtt{sinput}_{2} \right) \\
+1 - \tanh^{2}\left( \mathtt{W\_2}_{1} + \mathtt{W\_1}_{1,1} \mathtt{sinput}_{1} + \mathtt{W\_1}_{1,2} \mathtt{sinput}_{2} \right) \\
 \end{array}
 \right]
 \end{equation}
@@ -105,7 +105,7 @@ spb[1].L1.b |> latexify
 L"\begin{equation}
 \left[
 \begin{array}{c}
-1 - \tanh^{2}\left( \mathtt{b\_1}_{1} + \mathtt{W\_1}_{1,1} \mathtt{sinput}_{1} + \mathtt{W\_1}_{1,2} \mathtt{sinput}_{2} \right) \\
+1 - \tanh^{2}\left( \mathtt{W\_2}_{1} + \mathtt{W\_1}_{1,1} \mathtt{sinput}_{1} + \mathtt{W\_1}_{1,2} \mathtt{sinput}_{2} \right) \\
 \end{array}
 \right]
 \end{equation}
diff --git a/src/layers/abstract.jl b/src/layers/abstract.jl
index 05e34db..ab81396 100644
--- a/src/layers/abstract.jl
+++ b/src/layers/abstract.jl
@@ -6,11 +6,13 @@ Obtain the symbolic parameters of a neural network model.
 # Examples
 
 ```jldoctest
-using SymbolicNeuralNetworks
+using SymbolicNeuralNetworks: symbolize!
 using AbstractNeuralNetworks
 
+cache = Dict()
 d = Dense(4, 5, tanh)
-symbolicparameters(d) |> typeof
+params = NeuralNetwork(Chain(d)).params.L1
+symbolize!(cache, params, :X) |> typeof
 
 # output
 
diff --git a/src/symbolic_neuralnet/symbolic_neuralnet.jl b/src/symbolic_neuralnet/symbolic_neuralnet.jl
index 73ec9be..f7a4dc5 100644
--- a/src/symbolic_neuralnet/symbolic_neuralnet.jl
+++ b/src/symbolic_neuralnet/symbolic_neuralnet.jl
@@ -27,12 +27,17 @@ struct SymbolicNeuralNetwork{   AT,
     input::IT
 end
 
-function SymbolicNeuralNetwork(arch::Architecture, model::Model)
-    # Generation of symbolic paramters
-    sparams = symbolicparameters(model)
-    @variables sinput[1:input_dimension(model)]
+function SymbolicNeuralNetwork(nn::NeuralNetwork)
+    cache = Dict()
+    sparams = symbolize!(cache, nn.params, :W)
+    @variables sinput[1:input_dimension(nn.model)]
+
+    SymbolicNeuralNetwork(nn.architecture, nn.model, sparams, sinput)
+end
 
-    SymbolicNeuralNetwork(arch, model, sparams, sinput)
+function SymbolicNeuralNetwork(arch::Architecture, model::Model)
+    nn = NeuralNetwork(arch, model, CPU(), Float64)
+    SymbolicNeuralNetwork(nn)
 end
 
 function SymbolicNeuralNetwork(model::Chain)
diff --git a/src/symbolic_neuralnet/symbolize.jl b/src/symbolic_neuralnet/symbolize.jl
index 45b7dae..9738059 100644
--- a/src/symbolic_neuralnet/symbolize.jl
+++ b/src/symbolic_neuralnet/symbolize.jl
@@ -1,48 +1,132 @@
-#=
-    This files contains recursive functions to create a preserving shape symbolic params which can be the form of any combinaition of Tuple, namedTuple, Array and Real. 
-=#
+"""
+    symboliccounter!(cache, arg; redundancy)
 
-function SymbolicName(arg, storage; redundancy = true)
+Add a specific argument to the cache.
+
+# Examples
+
+```jldoctest
+using SymbolicNeuralNetworks: symboliccounter!
+
+cache = Dict()
+var = symboliccounter!(cache, :var)
+(cache, var)
+
+# output
+(Dict{Any, Any}(:var => 1), :var_1)
+
+```
+"""
+function symboliccounter!(cache::Dict, arg::Symbol; redundancy::Bool = true)
     if redundancy
-        arg ∈ keys(storage) ? storage[arg] += 1 : storage[arg] = 1
-        nam = string(arg)*"_"*string(storage[arg])
-        return Symbol(nam)
+        arg ∈ keys(cache) ? cache[arg] += 1 : cache[arg] = 1
+        nam = string(arg) * "_" * string(cache[arg])
+        Symbol(nam)
     else
-        nam = string(arg)
-        return Symbol(nam)
+        arg
     end
 end
 
-function symbolize(::Real, var_name::Union{Missing, Symbol} = missing, storage = Dict(); redundancy = true)
-    sname = ismissing(var_name) ?  SymbolicName(:X, storage; redundancy = redundancy) : SymbolicName(var_name, storage; redundancy = redundancy)
-    ((@variables $sname)[1], storage)
-end
+"""
+    symbolize!(cache, nt, var_name)
+
+Symbolize all the arguments in `nt`.
+
+# Examples
+
+```jldoctest
+using SymbolicNeuralNetworks: symbolize!
+
+cache = Dict()
+sym = symbolize!(cache, .1, :X)
+(sym, cache)
+
+# output
+
+(X_1, Dict{Any, Any}(:X => 1))
+```
+
+```jldoctest
+using SymbolicNeuralNetworks: symbolize!
+
+cache = Dict()
+arr = rand(2, 1)
+sym_scalar = symbolize!(cache, .1, :X)
+sym_array = symbolize!(cache, arr, :Y)
+(sym_array, cache)
+
+# output
+
+(Y_1[Base.OneTo(2),Base.OneTo(1)], Dict{Any, Any}(:X => 1, :Y => 1))
+```
+
+Note that the for the second case the cache is storing a scalar under `:X` and an array under `:Y`. If we use the same label for both we get:
+
+```jldoctest
+using SymbolicNeuralNetworks: symbolize!
+
+cache = Dict()
+arr = rand(2, 1)
+sym_scalar = symbolize!(cache, .1, :X)
+sym_array = symbolize!(cache, arr, :X)
+(sym_array, cache)
 
-function symbolize(M::AbstractArray, var_name::Union{Missing, Symbol} = missing, storage = Dict(); redundancy = true)
-    sname = ismissing(var_name) ?  SymbolicName(:M, storage; redundancy = redundancy) : SymbolicName(var_name, storage; redundancy = redundancy)
-    ((@variables $sname[Tuple([1:s for s in size(M)])...])[1], storage)
+# output
+
+(X_2[Base.OneTo(2),Base.OneTo(1)], Dict{Any, Any}(:X => 2))
+```
+
+We can also use `symbolize!` with `NamedTuple`s:
+
+```jldoctest
+using SymbolicNeuralNetworks: symbolize!
+
+cache = Dict()
+nt = (a = 1, b = [1, 2])
+sym = symbolize!(cache, nt, :X)
+(sym, cache)
+
+# output
+
+((a = X_1, b = X_2[Base.OneTo(2)]), Dict{Any, Any}(:X => 2))
+```
+
+And for neural network parameters:
+
+```jldoctest
+using SymbolicNeuralNetworks: symbolize!
+using AbstractNeuralNetworks
+
+nn = NeuralNetwork(Chain(Dense(1, 2; use_bias = false), Dense(2, 1; use_bias = false)))
+cache = Dict()
+sym = symbolize!(cache, nn.params, :X) |> typeof
+
+# output
+
+NeuralNetworkParameters{(:L1, :L2), Tuple{@NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}}, @NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}}}}
+```
+
+# Implementation
+
+Internally this is using [`symboliccounter!`](@ref). This function is also adjusting/altering the `cache` (that is optionally supplied as an input argument).
+"""
+symbolize!
+
+function symbolize!(cache::Dict, ::Real, var_name::Symbol; redundancy::Bool = true)::Symbolics.Num
+    sname = symboliccounter!(cache, var_name; redundancy = redundancy)
+    (@variables $sname)[1]
 end
 
-function symbolize(nt::NamedTuple, var_name::Union{Missing, Symbol} = missing, storage = Dict(); redundancy = true)
-    if length(nt) == 1
-        symb, storage= symbolize(values(nt)[1], keys(nt)[1], storage; redundancy = redundancy)
-        return NamedTuple{keys(nt)}((symb,)), storage
-    else
-        symb, storage = symbolize(values(nt)[1], keys(nt)[1], storage; redundancy = redundancy)
-        symbs, storage = symbolize(NamedTuple{keys(nt)[2:end]}(values(nt)[2:end]), var_name, storage; redundancy = redundancy)
-        return  (NamedTuple{keys(nt)}(Tuple([symb, symbs...])), storage)
-    end
+function symbolize!(cache::Dict, M::AbstractArray, var_name::Symbol; redundancy::Bool = true)
+    sname = symboliccounter!(cache, var_name; redundancy = redundancy)
+    (@variables $sname[axes(M)...])[1]
 end
 
-function symbolize(t::Tuple, var_name::Union{Missing, Symbol} = missing, storage = Dict(); redundancy = true)
-    if length(t) == 1
-        symb, storage = symbolize(t[1], var_name, storage; redundancy = redundancy)
-        return (symb,), storage
-    else
-        symb, storage = symbolize(t[1], var_name, storage; redundancy = redundancy)
-        symbs, storage = symbolize(t[2:end], var_name, storage; redundancy = redundancy)
-        return (Tuple([symb, symbs...]), storage)
-    end
+function symbolize!(cache::Dict, nt::NamedTuple, var_name::Symbol; redundancy::Bool = true)
+    values = Tuple(symbolize!(cache, nt[key], var_name; redundancy = redundancy) for key in keys(nt))
+    NamedTuple{keys(nt)}(values)
 end
 
-#symbolize(nn::NeuralNetwork) = symbolize(nn.params, missing, Dict())[1]
\ No newline at end of file
+function symbolize!(cache::Dict, nt::NeuralNetworkParameters, var_name::Symbol; redundancy::Bool = true)
+    NeuralNetworkParameters(symbolize!(cache, nt.params, var_name; redundancy = redundancy))
+end
\ No newline at end of file
diff --git a/test/derivatives/symbolic_gradient.jl b/test/derivatives/symbolic_gradient.jl
index fabe3ac..028d8c5 100644
--- a/test/derivatives/symbolic_gradient.jl
+++ b/test/derivatives/symbolic_gradient.jl
@@ -14,19 +14,18 @@ This test checks if we perform the parallelization in the correct way.
 function test_symbolic_gradient(input_dim::Integer = 3, output_dim::Integer = 1, hidden_dim::Integer = 2, T::DataType = Float64, second_dim::Integer = 3)
     @assert second_dim > 1 "second_dim must be greater than 1!"
     c = Chain(Dense(input_dim, hidden_dim, tanh), Dense(hidden_dim, output_dim, tanh))
-    sparams = symbolicparameters(c)
-    ps = initialparameters(c, T) |> NeuralNetworkParameters
-    @variables sinput[1:input_dim]
-    sout = norm(c(sinput, sparams)) ^ 2
-    sdparams = symbolic_differentials(sparams)
+    nn = NeuralNetwork(c)
+    snn = SymbolicNeuralNetwork(nn)
+    sout = norm(c(snn.input, snn.params)) ^ 2
+    sdparams = symbolic_differentials(snn.params)
     _sgrad = symbolic_derivative(sout, sdparams)
     input = rand(T, input_dim, second_dim)
     for k in 1:second_dim
-        zgrad = Zygote.gradient(ps -> (norm(c(input[:, k], ps)) ^ 2), ps)[1].params
+        zgrad = Zygote.gradient(ps -> (norm(c(input[:, k], ps)) ^ 2), nn.params)[1].params
         for key1 in keys(_sgrad)
             for key2 in keys(_sgrad[key1])
-                executable_gradient = _build_nn_function(_sgrad[key1][key2], sparams, sinput)
-                sgrad = executable_gradient(input, ps, k)
+                executable_gradient = _build_nn_function(_sgrad[key1][key2], snn.params, snn.input)
+                sgrad = executable_gradient(input, nn.params, k)
                 @test sgrad ≈ zgrad[key1][key2]
             end
         end
@@ -39,15 +38,14 @@ Also checks the parallelization, but for the full function.
 """
 function test_symbolic_gradient2(input_dim::Integer = 3, output_dim::Integer = 1, hidden_dim::Integer = 2, T::DataType = Float64, second_dim::Integer = 1, third_dim::Integer = 1)
     c = Chain(Dense(input_dim, hidden_dim, tanh), Dense(hidden_dim, output_dim, tanh))
-    sparams = symbolicparameters(c)
-    ps = initialparameters(c, T) |> NeuralNetworkParameters
-    @variables sinput[1:input_dim]
-    sout = norm(c(sinput, sparams)) ^ 2
+    nn = NeuralNetwork(c, T)
+    snn = SymbolicNeuralNetwork(nn)
+    sout = norm(c(snn.input, snn.params)) ^ 2
     input = rand(T, input_dim, second_dim, third_dim)
-    zgrad = Zygote.gradient(ps -> (norm(c(input, ps)) ^ 2), ps)[1].params
-    sdparams = symbolic_differentials(sparams)
+    zgrad = Zygote.gradient(ps -> (norm(c(input, ps)) ^ 2), nn.params)[1].params
+    sdparams = symbolic_differentials(snn.params)
     _sgrad = symbolic_derivative(sout, sdparams)
-    sgrad = build_nn_function(_sgrad, sparams, sinput)(input, ps)
+    sgrad = build_nn_function(_sgrad, sparams, sinput)(input, nn.params)
     for key1 in keys(sgrad) for key2 in keys(sgrad[key1]) @test zgrad[key1][key2] ≈ sgrad[key1][key2] end end
 end
 
diff --git a/test/symbolic_neuralnet/symbolize.jl b/test/symbolic_neuralnet/symbolize.jl
index d3bc141..36eac72 100644
--- a/test/symbolic_neuralnet/symbolize.jl
+++ b/test/symbolic_neuralnet/symbolize.jl
@@ -1,22 +1,27 @@
+using AbstractNeuralNetworks: NeuralNetworkParameters
 using SymbolicNeuralNetworks
+using SymbolicNeuralNetworks: symbolize!
 using Symbolics
 using Test
 
 # a tuple of `NamedTuple`s and `Tuple`s.
-params = (  (W = [1,1], b = [2, 2]),  
-            (W = (4 , [1,2]), c = 2), 
-            [4 5; 7 8], 
-            (7, 8.2)
-            )
+params = NeuralNetworkParameters(  
+        (L1 = (W = [1,1], b = [2, 2]),
+         L2 = (W = (a = 4 , b = [1,2]), c = 2),
+         L3 = [4 5; 7 8],
+         L4 = (a = 7, b = 8.2)
+        ) )
 
-sparams = symbolize(params)[1]
+cache = Dict()
+sparams = symbolize!(cache, params, :W)
 
-@variables W_1[1:2] W_2 W_3[1:2] b_1[1:2] c_1 M_1[1:2, 1:2] X_1 X_2
+@variables W_1[Base.OneTo(2)] W_3 W_4[Base.OneTo(2)] W_2[Base.OneTo(2)] W_5 W_6[Base.OneTo(2), Base.OneTo(2)] W_7 W_8
 
-verified_sparams = (    (W = W_1, b = b_1), 
-                        (W = (W_2 , W_3), c = c_1), 
-                        M_1, 
-                        (X_1, X_2)
-                        )
+verified_sparams = NeuralNetworkParameters(
+        (L1 = (W = W_1, b = W_2), 
+         L2 = (W = (a = W_3 , b = W_4), c = W_5), 
+         L3 = W_6, 
+         L4 = (a = W_7, b = W_8)
+         ) )
 
 @test sparams === verified_sparams
\ No newline at end of file

From 5a60aacfddf2ecbac85b2f7bdb7546b9897342a6 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Sat, 14 Dec 2024 21:06:44 +0100
Subject: [PATCH 12/23] Deleted extra specific layer constructors that aren't
 needed anymore.

---
 src/SymbolicNeuralNetworks.jl                |  6 ------
 src/chain/chain.jl                           |  5 -----
 src/layers/abstract.jl                       | 22 --------------------
 src/layers/dense.jl                          |  9 --------
 src/layers/linear.jl                         |  5 -----
 src/symbolic_neuralnet/symbolic_neuralnet.jl |  6 ++++--
 6 files changed, 4 insertions(+), 49 deletions(-)
 delete mode 100644 src/chain/chain.jl
 delete mode 100644 src/layers/abstract.jl
 delete mode 100644 src/layers/dense.jl
 delete mode 100644 src/layers/linear.jl

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index b07d4c0..f46d136 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -22,12 +22,6 @@ module SymbolicNeuralNetworks
     export AbstractSymbolicNeuralNetwork
     export SymbolicNeuralNetwork
 
-    # make symbolic parameters (`NeuralNetworkParameters`)
-    include("layers/abstract.jl")
-    include("layers/dense.jl")
-    include("layers/linear.jl")
-    include("chain/chain.jl")
-
     include("symbolic_neuralnet/symbolic_neuralnet.jl")
 
     export build_nn_function
diff --git a/src/chain/chain.jl b/src/chain/chain.jl
deleted file mode 100644
index 419f039..0000000
--- a/src/chain/chain.jl
+++ /dev/null
@@ -1,5 +0,0 @@
-function symbolicparameters(model::Chain)
-    vals = symbolize(Tuple(symbolicparameters(layer) for layer in model))[1]
-    keys = Tuple(Symbol("L$(i)") for i in 1:length(vals))
-    NeuralNetworkParameters(NamedTuple{keys}(vals))
-end
\ No newline at end of file
diff --git a/src/layers/abstract.jl b/src/layers/abstract.jl
deleted file mode 100644
index ab81396..0000000
--- a/src/layers/abstract.jl
+++ /dev/null
@@ -1,22 +0,0 @@
-"""
-    symbolicparameters(model)
-
-Obtain the symbolic parameters of a neural network model.
-
-# Examples
-
-```jldoctest
-using SymbolicNeuralNetworks: symbolize!
-using AbstractNeuralNetworks
-
-cache = Dict()
-d = Dense(4, 5, tanh)
-params = NeuralNetwork(Chain(d)).params.L1
-symbolize!(cache, params, :X) |> typeof
-
-# output
-
-@NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}, b::Symbolics.Arr{Symbolics.Num, 1}}
-```
-"""
-symbolicparameters(model::Model) = error("symbolicparameters not implemented for model type ", typeof(model))
\ No newline at end of file
diff --git a/src/layers/dense.jl b/src/layers/dense.jl
deleted file mode 100644
index a35a794..0000000
--- a/src/layers/dense.jl
+++ /dev/null
@@ -1,9 +0,0 @@
-function symbolicparameters(::Dense{M, N, true}) where {M,N}
-    @variables W[1:N, 1:M], b[1:N]
-    (W = W, b = b)
-end
-
-function symbolicparameters(::Dense{M, N, false}) where {M,N}
-    @variables W[1:N, 1:M]
-    (W = W,)
-end
\ No newline at end of file
diff --git a/src/layers/linear.jl b/src/layers/linear.jl
deleted file mode 100644
index c4f3203..0000000
--- a/src/layers/linear.jl
+++ /dev/null
@@ -1,5 +0,0 @@
-# this should be superfluous!
-# function symbolicparameters(::Linear{M, N}) where {M, N}
-#     @variables W[1:N, 1:M]
-#     (W = W, )
-# end
\ No newline at end of file
diff --git a/src/symbolic_neuralnet/symbolic_neuralnet.jl b/src/symbolic_neuralnet/symbolic_neuralnet.jl
index f7a4dc5..da5a2c0 100644
--- a/src/symbolic_neuralnet/symbolic_neuralnet.jl
+++ b/src/symbolic_neuralnet/symbolic_neuralnet.jl
@@ -5,6 +5,8 @@ abstract type AbstractSymbolicNeuralNetwork{AT} <: AbstractNeuralNetwork{AT} end
 
 A symbolic neural network realizes a symbolic represenation (of small neural networks).
 
+# Fields
+
 The `struct` has the following fields:
 - `architecture`: the neural network architecture,
 - `model`: the model (typically a Chain that is the realization of the architecture),
@@ -13,9 +15,9 @@ The `struct` has the following fields:
 
 # Constructors
 
-    SymbolicNeuralNetwork(arch)
+    SymbolicNeuralNetwork(nn)
 
-Make a `SymbolicNeuralNetwork` based on an architecture and a set of equations.
+Make a `SymbolicNeuralNetwork` based on a `AbstractNeuralNetworks.Network`.
 """
 struct SymbolicNeuralNetwork{   AT, 
                                 MT, 

From d26cf279131d20daa5c5a06d5783bc6c00518db9 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Sun, 15 Dec 2024 20:06:41 +0100
Subject: [PATCH 13/23] Renamed file build_function2.jl to
 build_function_double_input.jl.

---
 src/SymbolicNeuralNetworks.jl                                   | 2 +-
 .../{build_function2.jl => build_function_double_input.jl}      | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename src/build_function/{build_function2.jl => build_function_double_input.jl} (100%)

diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index f46d136..49ea427 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -26,7 +26,7 @@ module SymbolicNeuralNetworks
 
     export build_nn_function
     include("build_function/build_function.jl")
-    include("build_function/build_function2.jl")
+    include("build_function/build_function_double_input.jl")
     include("build_function/build_function_arrays.jl")
 
     export SymbolicPullback
diff --git a/src/build_function/build_function2.jl b/src/build_function/build_function_double_input.jl
similarity index 100%
rename from src/build_function/build_function2.jl
rename to src/build_function/build_function_double_input.jl

From 85aa3f0a85c706787094dcab4ef8e2d0c139fb1d Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Sun, 15 Dec 2024 21:45:41 +0100
Subject: [PATCH 14/23] Added tests for build_nn_function (with input and with
 input&output); partially from docstring tests.

---
 src/build_function/build_function.jl          | 14 ++++++-----
 test/build_function/build_function.jl         | 24 ++++++++++++++++++
 .../build_function_double_input.jl            | 25 +++++++++++++++++++
 test/runtests.jl                              |  3 ++-
 4 files changed, 59 insertions(+), 7 deletions(-)
 create mode 100644 test/build_function/build_function.jl
 create mode 100644 test/build_function/build_function_double_input.jl

diff --git a/src/build_function/build_function.jl b/src/build_function/build_function.jl
index 33d532f..c89eb20 100644
--- a/src/build_function/build_function.jl
+++ b/src/build_function/build_function.jl
@@ -39,24 +39,26 @@ Build a function that can process a matrix. This is used as a starting point for
 # Examples
 
 ```jldoctest
-using SymbolicNeuralNetworks: _build_nn_function
-using SymbolicNeuralNetworks
+using SymbolicNeuralNetworks: _build_nn_function, SymbolicNeuralNetwork
 using AbstractNeuralNetworks
+import Random
+Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
 nn = NeuralNetwork(c)
 snn = SymbolicNeuralNetwork(nn)
 eq = c(snn.input, snn.params)
 built_function = _build_nn_function(eq, snn.params, snn.input)
-input = rand(2, 2)
-
-(built_function(input, nn.params, 1), built_function(input, nn.params, 2)) .≈ (c(input[:, 1], nn.params), c(input[:, 2], nn.params))
+built_function([1. 2.; 3. 4.], nn.params, 1)
 
 # output
 
-(true, true)
+1-element Vector{Float64}:
+ -0.9999967113439513
 ```
 
+Note that we have to supply an extra argument (index) to `_build_nn_function` that we do not have to supply to [`build_nn_function`](@ref).
+
 # Implementation
 
 This first calls `Symbolics.build_function` with the keyword argument `expression = Val{true}` and then modifies the generated code by calling:
diff --git a/test/build_function/build_function.jl b/test/build_function/build_function.jl
new file mode 100644
index 0000000..2c2ea34
--- /dev/null
+++ b/test/build_function/build_function.jl
@@ -0,0 +1,24 @@
+using SymbolicNeuralNetworks: _build_nn_function
+using SymbolicNeuralNetworks
+using AbstractNeuralNetworks
+using Test
+
+# this tests the function '_build_nn_function' (not 'build_nn_function')
+function apply_build_function(input_dim::Integer=2, output_dim::Inetger=1, num_examples::Integer=3)
+    c = Chain(Dense(input_dim, output_dim, tanh))
+    nn = NeuralNetwork(c)
+    snn = SymbolicNeuralNetwork(nn)
+    eq = c(snn.input, snn.params)
+    built_function = _build_nn_function(eq, snn.params, snn.input)
+    input = rand(input_dim, num_examples)
+
+    @test all(i -> (built_function(input, nn.params, i) ≈ c(input[:, i], nn.params)), 1:num_examples)
+end
+
+for input_dim ∈ (2, 3)
+    for ouput_dim ∈ (1, 2)
+        for num_examples ∈ (1, 2, 3)
+            apply_build_function(input_dim, output_dim, num_examples)
+        end
+    end
+end
\ No newline at end of file
diff --git a/test/build_function/build_function_double_input.jl b/test/build_function/build_function_double_input.jl
new file mode 100644
index 0000000..3a78825
--- /dev/null
+++ b/test/build_function/build_function_double_input.jl
@@ -0,0 +1,25 @@
+using SymbolicNeuralNetworks: _build_nn_function
+using SymbolicNeuralNetworks
+using AbstractNeuralNetworks
+using Test
+
+# this tests the function '_build_nn_function' (not 'build_nn_function') for input and output
+function apply_build_function(input_dim::Integer=2, output_dim::Inetger=1, num_examples::Integer=3)
+    c = Chain(Dense(input_dim, output_dim, tanh))
+    nn = NeuralNetwork(c)
+    snn = SymbolicNeuralNetwork(nn)
+    eq = c(snn.input, snn.params)
+    built_function = _build_nn_function(eq, snn.params, snn.input)
+    input = rand(input_dim, num_examples)
+    output = rand(output_dim, num_examples)
+
+    @test all(i -> (built_function(input, output, nn.params, i) ≈ c(input[:, i], output[:, i], nn.params)), 1:num_examples)
+end
+
+for input_dim ∈ (2, 3)
+    for ouput_dim ∈ (1, 2)
+        for num_examples ∈ (1, 2, 3)
+            apply_build_function(input_dim, output_dim, num_examples)
+        end
+    end
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index c094b2b..f2864ea 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,4 +5,5 @@ using Test
 @safetestset "Docstring tests.                                                                       " begin include("doctest.jl") end
 @safetestset "Symbolic gradient                                                                      " begin include("derivatives/symbolic_gradient.jl") end
 @safetestset "Symbolic Neural network                                                                " begin include("derivatives/jacobian.jl") end
-@safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end
\ No newline at end of file
+@safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end
+@safetestset "Tests associated to 'build_function.jl'                                                " begin include("build_function/build_function.jl") end
\ No newline at end of file

From cc8f1dd150237af79f9e9d65c817a982f40d447f Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Sun, 15 Dec 2024 23:33:27 +0100
Subject: [PATCH 15/23] Moved test from docstring to
 test/.../build_function_arrays.jl.

---
 Project.toml                                  |  3 +-
 src/build_function/build_function_arrays.jl   | 23 +++++++--------
 test/build_function/build_function.jl         |  4 +--
 test/build_function/build_function_arrays.jl  | 28 +++++++++++++++++++
 .../build_function_double_input.jl            | 12 ++++----
 test/runtests.jl                              |  5 ++--
 6 files changed, 51 insertions(+), 24 deletions(-)
 create mode 100644 test/build_function/build_function_arrays.jl

diff --git a/Project.toml b/Project.toml
index 48d1c9a..080b343 100644
--- a/Project.toml
+++ b/Project.toml
@@ -16,7 +16,6 @@ Documenter = "1.8.0"
 ForwardDiff = "0.10.38"
 Latexify = "0.16.5"
 RuntimeGeneratedFunctions = "0.5"
-SafeTestsets = "0.1"
 Symbolics = "5, 6"
 Zygote = "0.6.73"
 julia = "1.6"
@@ -31,4 +30,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "ForwardDiff", "Random", "Documenter", "Latexify", "SafeTestsets", "Zygote"]
\ No newline at end of file
+test = ["Test", "ForwardDiff", "Random", "Documenter", "Latexify", "SafeTestsets", "Zygote"]
diff --git a/src/build_function/build_function_arrays.jl b/src/build_function/build_function_arrays.jl
index d53a6ee..601d0f1 100644
--- a/src/build_function/build_function_arrays.jl
+++ b/src/build_function/build_function_arrays.jl
@@ -1,32 +1,29 @@
 """
     build_nn_function(eqs::AbstractArray{<:NeuralNetworkParameters}, sparams, sinput...)
 
-Build an executable function based on `eqs` that potentially also has a symbolic output.
+Build an executable function based on an array of symbolic equations `eqs`.
 
 # Examples
 
 ```jldoctest
 using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense, initialparameters, NeuralNetworkParameters
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
 import Random
 Random.seed!(123)
 
 ch = Chain(Dense(2, 1, tanh))
-nn = SymbolicNeuralNetwork(ch)
-eqs = [(a = ch(nn.input, nn.params), b = ch(nn.input, nn.params).^2), (c = ch(nn.input, nn.params).^3, )]
-funcs = build_nn_function(eqs, nn.params, nn.input)
+nn = NeuralNetwork(ch)
+snn = SymbolicNeuralNetwork(nn)
+eqs = [(a = ch(snn.input, snn.params), b = ch(snn.input, snn.params).^2), (c = ch(snn.input, snn.params).^3, )]
+funcs = build_nn_function(eqs, snn.params, snn.input)
 input = [1., 2.]
-ps = initialparameters(ch) |> NeuralNetworkParameters
-a = ch(input, ps)
-b = ch(input, ps).^2
-c = ch(input, ps).^3
-funcs_evaluated = funcs(input, ps)
-
-(funcs_evaluated[1].a, funcs_evaluated[1].b, funcs_evaluated[2].c) .≈ (a, b, c)
+funcs_evaluated = funcs(input, nn.params)
 
 # output
 
-(true, true, true)
+2-element Vector{NamedTuple}:
+ (a = [-0.9999386280616135], b = [0.9998772598897417])
+ (c = [-0.9998158954841537],)
 ```
 """
 function build_nn_function(eqs::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...)
diff --git a/test/build_function/build_function.jl b/test/build_function/build_function.jl
index 2c2ea34..5aef4c6 100644
--- a/test/build_function/build_function.jl
+++ b/test/build_function/build_function.jl
@@ -4,7 +4,7 @@ using AbstractNeuralNetworks
 using Test
 
 # this tests the function '_build_nn_function' (not 'build_nn_function')
-function apply_build_function(input_dim::Integer=2, output_dim::Inetger=1, num_examples::Integer=3)
+function apply_build_function(input_dim::Integer=2, output_dim::Integer=1, num_examples::Integer=3)
     c = Chain(Dense(input_dim, output_dim, tanh))
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
@@ -16,7 +16,7 @@ function apply_build_function(input_dim::Integer=2, output_dim::Inetger=1, num_e
 end
 
 for input_dim ∈ (2, 3)
-    for ouput_dim ∈ (1, 2)
+    for output_dim ∈ (1, 2)
         for num_examples ∈ (1, 2, 3)
             apply_build_function(input_dim, output_dim, num_examples)
         end
diff --git a/test/build_function/build_function_arrays.jl b/test/build_function/build_function_arrays.jl
new file mode 100644
index 0000000..6582e70
--- /dev/null
+++ b/test/build_function/build_function_arrays.jl
@@ -0,0 +1,28 @@
+using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
+using AbstractNeuralNetworks: Chain, Dense, initialparameters, NeuralNetworkParameters
+using Test
+import Random
+Random.seed!(123)
+
+function build_function_for_array_valued_equation(input_dim::Integer=2, output_dim::Integer=1)
+    ch = Chain(Dense(input_dim, output_dim, tanh))
+    nn = NeuralNetwork(ch)
+    snn = SymbolicNeuralNetwork(nn)
+    eqs = [(a = ch(snn.input, snn.params), b = ch(snn.input, snn.params).^2), (c = ch(snn.input, snn.params).^3, )]
+    funcs = build_nn_function(eqs, nn.params, nn.input)
+    input = [1., 2.]
+    a = ch(input, ps)
+    b = ch(input, ps).^2
+    c = ch(input, ps).^3
+    funcs_evaluated = funcs(input, ps)
+    funcs_evaluated_as_vector = [funcs_evaluated[1].a, funcs_evaluated[1].b, funcs_evaluated[2].c]
+    result_of_standard_computation = [a, b, c]
+
+    @test funcs_evaluated_as_vector ≈ result_of_standard_computation
+end
+
+for input_dim ∈ (2, 3)
+    for output_dim ∈ (1, 2)
+        build_function_for_array_valued_equation(input_dim, output_dim)
+    end
+end
\ No newline at end of file
diff --git a/test/build_function/build_function_double_input.jl b/test/build_function/build_function_double_input.jl
index 3a78825..01854d8 100644
--- a/test/build_function/build_function_double_input.jl
+++ b/test/build_function/build_function_double_input.jl
@@ -1,23 +1,25 @@
 using SymbolicNeuralNetworks: _build_nn_function
 using SymbolicNeuralNetworks
+using Symbolics: @variables
 using AbstractNeuralNetworks
 using Test
 
 # this tests the function '_build_nn_function' (not 'build_nn_function') for input and output
-function apply_build_function(input_dim::Integer=2, output_dim::Inetger=1, num_examples::Integer=3)
+function apply_build_function(input_dim::Integer=2, output_dim::Integer=1, num_examples::Integer=3)
     c = Chain(Dense(input_dim, output_dim, tanh))
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
-    eq = c(snn.input, snn.params)
-    built_function = _build_nn_function(eq, snn.params, snn.input)
+    @variables soutput[1:output_dim]
+    eq = (c(snn.input, snn.params) - soutput).^2
+    built_function = _build_nn_function(eq, snn.params, snn.input, soutput)
     input = rand(input_dim, num_examples)
     output = rand(output_dim, num_examples)
 
-    @test all(i -> (built_function(input, output, nn.params, i) ≈ c(input[:, i], output[:, i], nn.params)), 1:num_examples)
+    @test all(i -> (built_function(input, output, nn.params, i) ≈ (c(input[:, i], nn.params) - output[:, i]).^2), 1:num_examples)
 end
 
 for input_dim ∈ (2, 3)
-    for ouput_dim ∈ (1, 2)
+    for output_dim ∈ (1, 2)
         for num_examples ∈ (1, 2, 3)
             apply_build_function(input_dim, output_dim, num_examples)
         end
diff --git a/test/runtests.jl b/test/runtests.jl
index f2864ea..417bdd9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,9 +1,10 @@
 using SymbolicNeuralNetworks
 using SafeTestsets
-using Test
 
 @safetestset "Docstring tests.                                                                       " begin include("doctest.jl") end
 @safetestset "Symbolic gradient                                                                      " begin include("derivatives/symbolic_gradient.jl") end
 @safetestset "Symbolic Neural network                                                                " begin include("derivatives/jacobian.jl") end
 @safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end
-@safetestset "Tests associated to 'build_function.jl'                                                " begin include("build_function/build_function.jl") end
\ No newline at end of file
+@safetestset "Tests associated with 'build_function.jl'                                              " begin include("build_function/build_function.jl") end
+@safetestset "Tests associated with 'build_function_double_input.jl'                                 " begin include("build_function/build_function_double_input.jl") end
+@safetestset "Tests associated with 'build_function_array.jl                                         " begin include("build_function/build_function_arrays.jl") end
\ No newline at end of file

From 4194d0ba90c6a52f1d78b4fc428a1de071e95dbf Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Sun, 15 Dec 2024 23:35:34 +0100
Subject: [PATCH 16/23] Now importing correct structs/functions.

Fixed typos.
---
 test/build_function/build_function_arrays.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/test/build_function/build_function_arrays.jl b/test/build_function/build_function_arrays.jl
index 6582e70..de5329a 100644
--- a/test/build_function/build_function_arrays.jl
+++ b/test/build_function/build_function_arrays.jl
@@ -1,5 +1,5 @@
 using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense, initialparameters, NeuralNetworkParameters
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
 using Test
 import Random
 Random.seed!(123)
@@ -9,12 +9,12 @@ function build_function_for_array_valued_equation(input_dim::Integer=2, output_d
     nn = NeuralNetwork(ch)
     snn = SymbolicNeuralNetwork(nn)
     eqs = [(a = ch(snn.input, snn.params), b = ch(snn.input, snn.params).^2), (c = ch(snn.input, snn.params).^3, )]
-    funcs = build_nn_function(eqs, nn.params, nn.input)
-    input = [1., 2.]
-    a = ch(input, ps)
-    b = ch(input, ps).^2
-    c = ch(input, ps).^3
-    funcs_evaluated = funcs(input, ps)
+    funcs = build_nn_function(eqs, snn.params, snn.input)
+    input = Vector(1:input_dim)
+    a = ch(input, nn.params)
+    b = ch(input, nn.params).^2
+    c = ch(input, nn.params).^3
+    funcs_evaluated = funcs(input, nn.params)
     funcs_evaluated_as_vector = [funcs_evaluated[1].a, funcs_evaluated[1].b, funcs_evaluated[2].c]
     result_of_standard_computation = [a, b, c]
 

From 02c67ea2981fd56b54cd7101ee2bd4d37ed21fc4 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Mon, 16 Dec 2024 09:55:09 +0100
Subject: [PATCH 17/23] Removed reference to HSNN.

---
 docs/src/hamiltonian_neural_network.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/hamiltonian_neural_network.md b/docs/src/hamiltonian_neural_network.md
index f3581c4..f92d412 100644
--- a/docs/src/hamiltonian_neural_network.md
+++ b/docs/src/hamiltonian_neural_network.md
@@ -35,7 +35,7 @@ z_data = randn(T, 2, n_points)
 nothing # hide
 ```
 
-We now specify a pullback [`HamiltonianSymbolicNeuralNetwork`](@ref):
+We now specify a pullback `HamiltonianSymbolicNeuralNetwork`
 
 ```julia hnn
 _pullback = SymbolicPullback(nn)

From f495dc83d7efe9519c64d1750f926653d9b4b8eb Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Mon, 16 Dec 2024 10:43:18 +0100
Subject: [PATCH 18/23] Moved remaining tests from docstring tests to
 build_function directory.

ps -> nn.params.

Vector can't be used on Tuple of Vectors (apparently).

[1., 2.] -> Vector(1:input_dim) (so that we can deal with flexible input dimensions.
---
 src/build_function/build_function_arrays.jl  | 18 ++++------
 test/build_function/build_function_arrays.jl | 38 +++++++++++++++++++-
 test/runtests.jl                             |  2 +-
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/src/build_function/build_function_arrays.jl b/src/build_function/build_function_arrays.jl
index 601d0f1..8536ed9 100644
--- a/src/build_function/build_function_arrays.jl
+++ b/src/build_function/build_function_arrays.jl
@@ -44,25 +44,21 @@ Return a function that takes an input, (optionally) an output and neural network
 
 ```jldoctest
 using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense, initialparameters, NeuralNetworkParameters
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
 import Random
 Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
-nn = SymbolicNeuralNetwork(c)
-eqs = (a = c(nn.input, nn.params), b = c(nn.input, nn.params).^2)
-funcs = build_nn_function(eqs, nn.params, nn.input)
+nn = NeuralNetwork(c)
+snn = SymbolicNeuralNetwork(nn)
+eqs = (a = c(snn.input, snn.params), b = c(snn.input, snn.params).^2)
+funcs = build_nn_function(eqs, snn.params, snn.input)
 input = [1., 2.]
-ps = initialparameters(c) |> NeuralNetworkParameters
-a = c(input, ps)
-b = c(input, ps).^2
-funcs_evaluated = funcs(input, ps)
-
-(funcs_evaluated.a, funcs_evaluated.b) .≈ (a, b)
+funcs_evaluated = funcs(input, nn.params)
 
 # output
 
-(true, true)
+(a = [-0.9999386280616135], b = [0.9998772598897417])
 ```
 
 # Implementation
diff --git a/test/build_function/build_function_arrays.jl b/test/build_function/build_function_arrays.jl
index de5329a..e9f0980 100644
--- a/test/build_function/build_function_arrays.jl
+++ b/test/build_function/build_function_arrays.jl
@@ -1,4 +1,4 @@
-using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
+using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork, function_valued_parameters
 using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
 using Test
 import Random
@@ -21,8 +21,44 @@ function build_function_for_array_valued_equation(input_dim::Integer=2, output_d
     @test funcs_evaluated_as_vector ≈ result_of_standard_computation
 end
 
+function build_function_for_named_tuple(input_dim::Integer=2, output_dim::Integer=1)
+    c = Chain(Dense(input_dim, output_dim, tanh))
+    nn = NeuralNetwork(c)
+    snn = SymbolicNeuralNetwork(nn)
+    eqs = (a = c(snn.input, snn.params), b = c(snn.input, snn.params).^2)
+    funcs = build_nn_function(eqs, snn.params, snn.input)
+    input = Vector(1:input_dim)
+    a = c(input, nn.params)
+    b = c(input, nn.params).^2
+    funcs_evaluated = funcs(input, nn.params)
+
+    funcs_evaluated_as_vector = [funcs_evaluated.a, funcs_evaluated.b]
+    result_of_standard_computation = [a, b]
+
+    @test funcs_evaluated_as_vector ≈ result_of_standard_computation
+end
+
+function function_valued_parameters_for_named_tuple(input_dim::Integer=2, output_dim::Integer=1)
+    c = Chain(Dense(input_dim, output_dim, tanh))
+    nn = NeuralNetwork(c)
+    snn = SymbolicNeuralNetwork(nn)
+    eqs = (a = c(snn.input, snn.params), b = c(snn.input, snn.params).^2)
+    funcs = function_valued_parameters(eqs, snn.params, snn.input)
+    input = Vector(1:input_dim)
+    a = c(input, nn.params)
+    b = c(input, nn.params).^2
+
+    funcs_evaluated_as_vector = [funcs.a(input, nn.params), funcs.b(input, nn.params)]
+    result_of_standard_computation = [a, b]
+    
+    @test funcs_evaluated_as_vector ≈ result_of_standard_computation
+end
+
+# we test in the following order: `function_valued_parameters` → `build_function` (for `NamedTuple`) → `build_function` (for `Array` of `NamedTuple`s) as this is also how the functions are built.
 for input_dim ∈ (2, 3)
     for output_dim ∈ (1, 2)
+        function_valued_parameters_for_named_tuple(input_dim, output_dim)
+        build_function_for_named_tuple(input_dim, output_dim)
         build_function_for_array_valued_equation(input_dim, output_dim)
     end
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 417bdd9..8a79390 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,7 +1,7 @@
 using SymbolicNeuralNetworks
 using SafeTestsets
 
-@safetestset "Docstring tests.                                                                       " begin include("doctest.jl") end
+# @safetestset "Docstring tests.                                                                       " begin include("doctest.jl") end
 @safetestset "Symbolic gradient                                                                      " begin include("derivatives/symbolic_gradient.jl") end
 @safetestset "Symbolic Neural network                                                                " begin include("derivatives/jacobian.jl") end
 @safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end

From a260ae5f5a9b8b91f1854f9b8531cd3bd072a9f0 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Mon, 16 Dec 2024 11:34:28 +0100
Subject: [PATCH 19/23] Added test that compares symbolic pullback to zygote
 pullback.

Flipped order of functions.

Added _get_contents method for Tuple as argument.

Added another method to deal with Zygote idiosyncracies.

Fixed method for _get_params.
---
 Project.toml                 |  4 +++-
 src/derivatives/pullback.jl  | 19 ++++++++++--------
 test/derivatives/pullback.jl | 38 ++++++++++++++++++++++++++++++++++++
 test/runtests.jl             |  3 ++-
 4 files changed, 54 insertions(+), 10 deletions(-)
 create mode 100644 test/derivatives/pullback.jl

diff --git a/Project.toml b/Project.toml
index 080b343..22fb6ef 100644
--- a/Project.toml
+++ b/Project.toml
@@ -14,6 +14,7 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 AbstractNeuralNetworks = "0.3, 0.4"
 Documenter = "1.8.0"
 ForwardDiff = "0.10.38"
+GeometricMachineLearning = "0.3.7"
 Latexify = "0.16.5"
 RuntimeGeneratedFunctions = "0.5"
 Symbolics = "5, 6"
@@ -28,6 +29,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+GeometricMachineLearning = "194d25b2-d3f5-49f0-af24-c124f4aa80cc"
 
 [targets]
-test = ["Test", "ForwardDiff", "Random", "Documenter", "Latexify", "SafeTestsets", "Zygote"]
+test = ["Test", "ForwardDiff", "Random", "Documenter", "Latexify", "SafeTestsets", "Zygote", "GeometricMachineLearning"]
diff --git a/src/derivatives/pullback.jl b/src/derivatives/pullback.jl
index dba9c53..b1edbdb 100644
--- a/src/derivatives/pullback.jl
+++ b/src/derivatives/pullback.jl
@@ -16,7 +16,7 @@ nn = SymbolicNeuralNetwork(c)
 loss = FeedForwardLoss()
 pb = SymbolicPullback(nn, loss)
 ps = initialparameters(c) |> NeuralNetworkParameters
-pv_values = pb(ps, nn.model, (rand(2), rand(1)))[2](1) |> typeof
+pb_values = pb(ps, nn.model, (rand(2), rand(1)))[2](1) |> typeof
 
 # output
 
@@ -47,19 +47,20 @@ import Random
 Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
-nn = SymbolicNeuralNetwork(c)
+nn = NeuralNetwork(c)
+snn = SymbolicNeuralNetwork(nn)
 loss = FeedForwardLoss()
-pb = SymbolicPullback(nn, loss)
-ps = initialparameters(c) |> NeuralNetworkParameters
+pb = SymbolicPullback(snn, loss)
 input_output = (rand(2), rand(1))
-loss_and_pullback = pb(ps, nn.model, input_output)
-pv_values = loss_and_pullback[2](1)
+loss_and_pullback = pb(nn.params, nn.model, input_output)
+# note that we apply the second argument to another input `1`
+pb_values = loss_and_pullback[2](1)
 
 @variables soutput[1:SymbolicNeuralNetworks.output_dimension(nn.model)]
 symbolic_pullbacks = SymbolicNeuralNetworks.symbolic_pullback(loss(nn.model, nn.params, nn.input, soutput), nn)
-pv_values2 = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput)(input_output[1], input_output[2], ps)
+pb_values2 = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput)(input_output[1], input_output[2], ps)
 
-pv_values == (pv_values2 |> SymbolicNeuralNetworks._get_params |> SymbolicNeuralNetworks._get_contents)
+pb_values == (pb_values2 |> SymbolicNeuralNetworks._get_params |> SymbolicNeuralNetworks._get_contents)
 
 # output
 
@@ -106,6 +107,7 @@ Return the `NamedTuple` that's equivalent to the `NeuralNetworkParameters`.
 """
 _get_params(nt::NamedTuple) = nt
 _get_params(ps::NeuralNetworkParameters) = ps.params
+_get_params(ps::NamedTuple{(:params,), Tuple{NT}}) where {NT<:NamedTuple} = ps.params
 _get_params(ps::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}) = [_get_params(nt) for nt in ps]
 
 """
@@ -134,6 +136,7 @@ function __get_contents(nt::AbstractArray{<:NamedTuple})
     nt
 end
 _get_contents(nt::AbstractArray{<:NamedTuple}) = __get_contents(nt)
+_get_contents(nt::Tuple{<:NamedTuple}) = nt[1]
 
 # (_pullback::SymbolicPullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps)
 function (_pullback::SymbolicPullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple
diff --git a/test/derivatives/pullback.jl b/test/derivatives/pullback.jl
new file mode 100644
index 0000000..9eee5c7
--- /dev/null
+++ b/test/derivatives/pullback.jl
@@ -0,0 +1,38 @@
+using SymbolicNeuralNetworks
+using SymbolicNeuralNetworks: _get_params, _get_contents
+using AbstractNeuralNetworks
+using Symbolics
+using GeometricMachineLearning: ZygotePullback
+using Test
+import Random
+Random.seed!(123)
+
+compare_values(arr1::Array, arr2::Array) = @test arr1 ≈ arr2
+function compare_values(nt1::NamedTuple, nt2::NamedTuple)
+    @assert keys(nt1) == keys(nt2)
+    NamedTuple{keys(nt1)}((compare_values(arr1, arr2) for (arr1, arr2) in zip(values(nt1), values(nt2))))
+end
+
+function compare_symbolic_pullback_to_zygote_pullback(input_dim::Integer, output_dim::Integer, second_dim::Integer=1)
+    c = Chain(Dense(input_dim, output_dim, tanh))
+    nn = NeuralNetwork(c)
+    snn = SymbolicNeuralNetwork(nn)
+    loss = FeedForwardLoss()
+    spb = SymbolicPullback(snn, loss)
+    input_output = (rand(input_dim, second_dim), rand(output_dim, second_dim))
+    loss_and_pullback = spb(nn.params, nn.model, input_output)
+    # note that we apply the second argument to another input `1`
+    pb_values = loss_and_pullback[2](1)
+
+    zpb = ZygotePullback(loss)
+    loss_and_pullback_zygote = zpb(nn.params, nn.model, input_output)
+    pb_values_zygote = loss_and_pullback_zygote[2](1) |> _get_contents |> _get_params
+
+    compare_values(pb_values, pb_values_zygote)
+end
+
+for input_dim ∈ (2, 3)
+    for output_dim ∈ (1, 2)
+        compare_symbolic_pullback_to_zygote_pullback(input_dim, output_dim)
+    end
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 8a79390..edb44a1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,4 +7,5 @@ using SafeTestsets
 @safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end
 @safetestset "Tests associated with 'build_function.jl'                                              " begin include("build_function/build_function.jl") end
 @safetestset "Tests associated with 'build_function_double_input.jl'                                 " begin include("build_function/build_function_double_input.jl") end
-@safetestset "Tests associated with 'build_function_array.jl                                         " begin include("build_function/build_function_arrays.jl") end
\ No newline at end of file
+@safetestset "Tests associated with 'build_function_array.jl                                         " begin include("build_function/build_function_arrays.jl") end
+@safetestset "Compare Zygote Pullback with Symbolic Pullback                                         " begin include("derivatives/pullback.jl") end
\ No newline at end of file

From 30a00456c1ecc2e1cf734d5c760e33eadd83fef6 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Mon, 16 Dec 2024 13:28:14 +0100
Subject: [PATCH 20/23] Fixed docstring problem.

---
 src/derivatives/pullback.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/derivatives/pullback.jl b/src/derivatives/pullback.jl
index b1edbdb..1c1c894 100644
--- a/src/derivatives/pullback.jl
+++ b/src/derivatives/pullback.jl
@@ -57,10 +57,10 @@ loss_and_pullback = pb(nn.params, nn.model, input_output)
 pb_values = loss_and_pullback[2](1)
 
 @variables soutput[1:SymbolicNeuralNetworks.output_dimension(nn.model)]
-symbolic_pullbacks = SymbolicNeuralNetworks.symbolic_pullback(loss(nn.model, nn.params, nn.input, soutput), nn)
-pb_values2 = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput)(input_output[1], input_output[2], ps)
+symbolic_pullbacks = SymbolicNeuralNetworks.symbolic_pullback(loss(nn.model, snn.params, snn.input, soutput), snn)
+pb_values2 = build_nn_function(symbolic_pullbacks, snn.params, snn.input, soutput)(input_output[1], input_output[2], nn.params)
 
-pb_values == (pb_values2 |> SymbolicNeuralNetworks._get_params |> SymbolicNeuralNetworks._get_contents)
+pb_values == (pb_values2 |> SymbolicNeuralNetworks._get_contents |> SymbolicNeuralNetworks._get_params)
 
 # output
 
@@ -128,15 +128,15 @@ _get_contents([(a = "element_contained_in_vector", )])
 ```
 """
 _get_contents(nt::NamedTuple) = nt
-function _get_contents(nt::AbstractVector{<:NamedTuple})
+function _get_contents(nt::AbstractVector{<:Union{NamedTuple, NeuralNetworkParameters}})
     length(nt) == 1 ? nt[1] : __get_contents(nt)
 end
-function __get_contents(nt::AbstractArray{<:NamedTuple})
+function __get_contents(nt::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}})
     @warn "The pullback returns an array expression. There is probably a bug in the code somewhere."
     nt
 end
-_get_contents(nt::AbstractArray{<:NamedTuple}) = __get_contents(nt)
-_get_contents(nt::Tuple{<:NamedTuple}) = nt[1]
+_get_contents(nt::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}) = __get_contents(nt)
+_get_contents(nt::Tuple{<:Union{NamedTuple, NeuralNetworkParameters}}) = nt[1]
 
 # (_pullback::SymbolicPullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps)
 function (_pullback::SymbolicPullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple

From 53643c58123307c52735b79ac641a9e68d435046 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Mon, 16 Dec 2024 15:27:14 +0100
Subject: [PATCH 21/23] Remove doctests completely.

---
 test/runtests.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index edb44a1..200da87 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,7 +1,6 @@
 using SymbolicNeuralNetworks
 using SafeTestsets
 
-# @safetestset "Docstring tests.                                                                       " begin include("doctest.jl") end
 @safetestset "Symbolic gradient                                                                      " begin include("derivatives/symbolic_gradient.jl") end
 @safetestset "Symbolic Neural network                                                                " begin include("derivatives/jacobian.jl") end
 @safetestset "Symbolic Params                                                                        " begin include("symbolic_neuralnet/symbolize.jl") end

From fe4dfa7a9b7c8d0ca2c34200b8990018e8cfd4fe Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Tue, 17 Dec 2024 10:03:04 +0100
Subject: [PATCH 22/23] params now used as function instead of as a keyword.

---
 docs/src/double_derivative.md                 |  4 +--
 docs/src/symbolic_neural_networks.md          |  6 ++--
 scripts/pullback_comparison.jl                | 12 +++----
 src/build_function/build_function.jl          | 10 +++---
 src/build_function/build_function_arrays.jl   | 22 ++++++------
 .../build_function_double_input.jl            |  2 +-
 src/derivatives/gradient.jl                   |  7 ++--
 src/derivatives/jacobian.jl                   |  4 +--
 src/derivatives/pullback.jl                   | 14 ++++----
 src/symbolic_neuralnet/symbolic_neuralnet.jl  |  6 ++--
 src/symbolic_neuralnet/symbolize.jl           |  8 ++---
 test/build_function/build_function.jl         |  7 ++--
 test/build_function/build_function_arrays.jl  | 34 +++++++++----------
 .../build_function_double_input.jl            |  7 ++--
 test/derivatives/pullback.jl                  |  5 +--
 test/derivatives/symbolic_gradient.jl         | 20 +++++------
 16 files changed, 87 insertions(+), 81 deletions(-)

diff --git a/docs/src/double_derivative.md b/docs/src/double_derivative.md
index b3f3c4b..46a8e4e 100644
--- a/docs/src/double_derivative.md
+++ b/docs/src/double_derivative.md
@@ -18,7 +18,7 @@
 ```@example jacobian_gradient
 using AbstractNeuralNetworks
 using SymbolicNeuralNetworks
-using SymbolicNeuralNetworks: Jacobian, Gradient, derivative
+using SymbolicNeuralNetworks: Jacobian, Gradient, derivative, params
 using Latexify: latexify
 
 c = Chain(Dense(2, 1, tanh; use_bias = false))
@@ -92,7 +92,7 @@ x = \begin{pmatrix} 1 \\ 0 \end{pmatrix}, \quad W = \begin{bmatrix} 1 & 0 \\ 0 &
 ```
 
 ```@example jacobian_gradient
-built_function = build_nn_function(derivative(g), nn.params, nn.input)
+built_function = build_nn_function(derivative(g), params(nn), nn.input)
 
 x = [1., 0.]
 ps = NeuralNetworkParameters((L1 = (W = [1. 0.; 0. 1.], b = [0., 0.]), ))
diff --git a/docs/src/symbolic_neural_networks.md b/docs/src/symbolic_neural_networks.md
index 027db31..7cb3b1c 100644
--- a/docs/src/symbolic_neural_networks.md
+++ b/docs/src/symbolic_neural_networks.md
@@ -6,7 +6,7 @@ We first call the symbolic neural network that only consists of one layer:
 
 ```@example snn
 using SymbolicNeuralNetworks
-using AbstractNeuralNetworks: Chain, Dense, initialparameters
+using AbstractNeuralNetworks: Chain, Dense, initialparameters, params
 
 input_dim = 2
 output_dim = 1
@@ -23,7 +23,7 @@ using Symbolics
 using Latexify: latexify
 
 @variables sinput[1:input_dim]
-soutput = nn.model(sinput, nn.params)
+soutput = nn.model(sinput, params(nn))
 
 soutput
 ```
@@ -101,7 +101,7 @@ We now compare the neural network-approximated curve to the original one:
 fig = Figure()
 ax = Axis3(fig[1, 1])
 
-surface!(x_vec, y_vec, [c([x, y], nn_cpu.params)[1] for x in x_vec, y in y_vec]; alpha = .8, colormap = :darkterrain, transparency = true)
+surface!(x_vec, y_vec, [c([x, y], params(nn_cpu))[1] for x in x_vec, y in y_vec]; alpha = .8, colormap = :darkterrain, transparency = true)
 fig
 ```
 
diff --git a/scripts/pullback_comparison.jl b/scripts/pullback_comparison.jl
index c9acee3..05b5d31 100644
--- a/scripts/pullback_comparison.jl
+++ b/scripts/pullback_comparison.jl
@@ -19,10 +19,10 @@ output = rand(1, batch_size)
 # output sensitivities
 _do = 1.
 
-# spb(nn_cpu.params, nn.model, (input, output))[2](_do)
-# zpb(nn_cpu.params, nn.model, (input, output))[2](_do)
-# @time spb_evaluated = spb(nn_cpu.params, nn.model, (input, output))[2](_do)
-# @time zpb_evaluated = zpb(nn_cpu.params, nn.model, (input, output))[2](_do)[1].params
+# spb(params(nn_cpu), nn.model, (input, output))[2](_do)
+# zpb(params(nn_cpu), nn.model, (input, output))[2](_do)
+# @time spb_evaluated = spb(params(nn_cpu), nn.model, (input, output))[2](_do)
+# @time zpb_evaluated = zpb(params(nn_cpu), nn.model, (input, output))[2](_do)[1].params
 # @assert values(spb_evaluated) .≈ values(zpb_evaluated)
 
 function timenn(pb, params, model, input, output, _do = 1.)
@@ -30,5 +30,5 @@ function timenn(pb, params, model, input, output, _do = 1.)
     @time pb(params, model, (input, output))[2](_do)
 end
 
-timenn(spb, nn_cpu.params, nn.model, input, output)
-timenn(zpb, nn_cpu.params, nn.model, input, output)
+timenn(spb, params(nn_cpu), nn.model, input, output)
+timenn(zpb, params(nn_cpu), nn.model, input, output)
diff --git a/src/build_function/build_function.jl b/src/build_function/build_function.jl
index c89eb20..4943989 100644
--- a/src/build_function/build_function.jl
+++ b/src/build_function/build_function.jl
@@ -19,7 +19,7 @@ The functions mentioned in the implementation section were adjusted ad-hoc to de
 Other problems may occur. In case you bump into one please [open an issue on github](https://github.com/JuliaGNI/SymbolicNeuralNetworks.jl/issues).
 """
 function build_nn_function(eq::EqT, nn::AbstractSymbolicNeuralNetwork)
-    build_nn_function(eq, nn.params, nn.input)
+    build_nn_function(eq, params(nn), nn.input)
 end
 
 function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr)
@@ -40,16 +40,16 @@ Build a function that can process a matrix. This is used as a starting point for
 
 ```jldoctest
 using SymbolicNeuralNetworks: _build_nn_function, SymbolicNeuralNetwork
-using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params, Chain, Dense, NeuralNetwork
 import Random
 Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
 nn = NeuralNetwork(c)
 snn = SymbolicNeuralNetwork(nn)
-eq = c(snn.input, snn.params)
-built_function = _build_nn_function(eq, snn.params, snn.input)
-built_function([1. 2.; 3. 4.], nn.params, 1)
+eq = c(snn.input, params(snn))
+built_function = _build_nn_function(eq, params(snn), snn.input)
+built_function([1. 2.; 3. 4.], params(nn), 1)
 
 # output
 
diff --git a/src/build_function/build_function_arrays.jl b/src/build_function/build_function_arrays.jl
index 8536ed9..c30737a 100644
--- a/src/build_function/build_function_arrays.jl
+++ b/src/build_function/build_function_arrays.jl
@@ -7,17 +7,17 @@ Build an executable function based on an array of symbolic equations `eqs`.
 
 ```jldoctest
 using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
 import Random
 Random.seed!(123)
 
 ch = Chain(Dense(2, 1, tanh))
 nn = NeuralNetwork(ch)
 snn = SymbolicNeuralNetwork(nn)
-eqs = [(a = ch(snn.input, snn.params), b = ch(snn.input, snn.params).^2), (c = ch(snn.input, snn.params).^3, )]
-funcs = build_nn_function(eqs, snn.params, snn.input)
+eqs = [(a = ch(snn.input, params(snn)), b = ch(snn.input, params(snn)).^2), (c = ch(snn.input, params(snn)).^3, )]
+funcs = build_nn_function(eqs, params(snn), snn.input)
 input = [1., 2.]
-funcs_evaluated = funcs(input, nn.params)
+funcs_evaluated = funcs(input, params(nn))
 
 # output
 
@@ -44,17 +44,17 @@ Return a function that takes an input, (optionally) an output and neural network
 
 ```jldoctest
 using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
 import Random
 Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
 nn = NeuralNetwork(c)
 snn = SymbolicNeuralNetwork(nn)
-eqs = (a = c(snn.input, snn.params), b = c(snn.input, snn.params).^2)
-funcs = build_nn_function(eqs, snn.params, snn.input)
+eqs = (a = c(snn.input, params(snn)), b = c(snn.input, params(snn)).^2)
+funcs = build_nn_function(eqs, params(snn), snn.input)
 input = [1., 2.]
-funcs_evaluated = funcs(input, nn.params)
+funcs_evaluated = funcs(input, params(nn))
 
 # output
 
@@ -83,14 +83,14 @@ Return an executable function for each entry in `eqs`. This still has to be proc
 
 ```jldoctest
 using SymbolicNeuralNetworks: function_valued_parameters, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense, initialparameters, NeuralNetworkParameters
+using AbstractNeuralNetworks: Chain, Dense, initialparameters, NeuralNetworkParameters, params
 import Random
 Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
 nn = SymbolicNeuralNetwork(c)
-eqs = (a = c(nn.input, nn.params), b = c(nn.input, nn.params).^2)
-funcs = function_valued_parameters(eqs, nn.params, nn.input)
+eqs = (a = c(nn.input, params(nn)), b = c(nn.input, params(nn)).^2)
+funcs = function_valued_parameters(eqs, params(nn), nn.input)
 input = [1., 2.]
 ps = initialparameters(c) |> NeuralNetworkParameters
 a = c(input, ps)
diff --git a/src/build_function/build_function_double_input.jl b/src/build_function/build_function_double_input.jl
index 77e1853..ece85c7 100644
--- a/src/build_function/build_function_double_input.jl
+++ b/src/build_function/build_function_double_input.jl
@@ -13,7 +13,7 @@ Also compare this to [`build_nn_function(::EqT, ::AbstractSymbolicNeuralNetwork)
 See the *extended help section* of [`build_nn_function(::EqT, ::AbstractSymbolicNeuralNetwork)`](@ref).
 """
 function build_nn_function(eqs, nn::AbstractSymbolicNeuralNetwork, soutput)
-    build_nn_function(eqs, nn.params, nn.input, soutput)
+    build_nn_function(eqs, params(nn), nn.input, soutput)
 end
 
 function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr)
diff --git a/src/derivatives/gradient.jl b/src/derivatives/gradient.jl
index ea519ec..8790a80 100644
--- a/src/derivatives/gradient.jl
+++ b/src/derivatives/gradient.jl
@@ -75,7 +75,7 @@ function Gradient(output::EqT, nn::SymbolicNeuralNetwork)
 end
 
 function Gradient(nn::SymbolicNeuralNetwork)
-    Gradient(nn.model(nn.input, nn.params), nn)
+    Gradient(nn.model(nn.input, params(nn)), nn)
 end
 
 @doc raw"""
@@ -90,12 +90,13 @@ This is used by [`Gradient`](@ref) and [`SymbolicPullback`](@ref).
 ```jldoctest
 using SymbolicNeuralNetworks: SymbolicNeuralNetwork, symbolic_pullback
 using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params
 using LinearAlgebra: norm
 using Latexify: latexify
 
 c = Chain(Dense(2, 1, tanh))
 nn = SymbolicNeuralNetwork(c)
-output = c(nn.input, nn.params)
+output = c(nn.input, params(nn))
 spb = symbolic_pullback(output, nn)
 
 spb[1].L1.b |> latexify
@@ -113,6 +114,6 @@ L"\begin{equation}
 ```
 """
 function symbolic_pullback(soutput::EqT, nn::AbstractSymbolicNeuralNetwork)::Union{AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, Union{NamedTuple, NeuralNetworkParameters}}
-    symbolic_diffs = symbolic_differentials(nn.params)
+    symbolic_diffs = symbolic_differentials(params(nn))
     [symbolic_derivative(soutput_single, symbolic_diffs) for soutput_single ∈ soutput]
 end
\ No newline at end of file
diff --git a/src/derivatives/jacobian.jl b/src/derivatives/jacobian.jl
index 5db7a7b..a64f0fe 100644
--- a/src/derivatives/jacobian.jl
+++ b/src/derivatives/jacobian.jl
@@ -18,7 +18,7 @@ The output of `Jacobian` consists of a `NamedTuple` that has the following keys:
 If `output` is not supplied as an input argument than it is taken to be:
 
 ```julia 
-soutput = nn.model(nn.input, nn.params)
+soutput = nn.model(nn.input, params(nn))
 ```
 
 # Implementation
@@ -82,7 +82,7 @@ derivative(j::Jacobian) = j.□
 function Jacobian(nn::AbstractSymbolicNeuralNetwork)
     
     # Evaluation of the symbolic output
-    soutput = nn.model(nn.input, nn.params)
+    soutput = nn.model(nn.input, params(nn))
 
     Jacobian(soutput, nn)
 end
diff --git a/src/derivatives/pullback.jl b/src/derivatives/pullback.jl
index 1c1c894..3e43153 100644
--- a/src/derivatives/pullback.jl
+++ b/src/derivatives/pullback.jl
@@ -41,7 +41,7 @@ We note the following seeming peculiarity:
 
 ```jldoctest
 using SymbolicNeuralNetworks
-using AbstractNeuralNetworks
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, FeedForwardLoss, params
 using Symbolics
 import Random
 Random.seed!(123)
@@ -52,13 +52,13 @@ snn = SymbolicNeuralNetwork(nn)
 loss = FeedForwardLoss()
 pb = SymbolicPullback(snn, loss)
 input_output = (rand(2), rand(1))
-loss_and_pullback = pb(nn.params, nn.model, input_output)
+loss_and_pullback = pb(params(nn), nn.model, input_output)
 # note that we apply the second argument to another input `1`
 pb_values = loss_and_pullback[2](1)
 
 @variables soutput[1:SymbolicNeuralNetworks.output_dimension(nn.model)]
-symbolic_pullbacks = SymbolicNeuralNetworks.symbolic_pullback(loss(nn.model, snn.params, snn.input, soutput), snn)
-pb_values2 = build_nn_function(symbolic_pullbacks, snn.params, snn.input, soutput)(input_output[1], input_output[2], nn.params)
+symbolic_pullbacks = SymbolicNeuralNetworks.symbolic_pullback(loss(nn.model, params(snn), snn.input, soutput), snn)
+pb_values2 = build_nn_function(symbolic_pullbacks, params(snn), snn.input, soutput)(input_output[1], input_output[2], params(nn))
 
 pb_values == (pb_values2 |> SymbolicNeuralNetworks._get_contents |> SymbolicNeuralNetworks._get_params)
 
@@ -88,9 +88,9 @@ end
 
 function SymbolicPullback(nn::SymbolicNeuralNetwork, loss::NetworkLoss)
     @variables soutput[1:output_dimension(nn.model)]
-    symbolic_loss = loss(nn.model, nn.params, nn.input, soutput)
+    symbolic_loss = loss(nn.model, params(nn), 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, params(nn), nn.input, soutput)
     function pbs(input, output, params)
         pullback(::Union{Real, AbstractArray{<:Real}}) = _get_contents(_get_params(pbs_executable(input, output, params)))
         pullback
@@ -106,7 +106,7 @@ SymbolicPullback(nn::SymbolicNeuralNetwork) = SymbolicPullback(nn, AbstractNeura
 Return the `NamedTuple` that's equivalent to the `NeuralNetworkParameters`.
 """
 _get_params(nt::NamedTuple) = nt
-_get_params(ps::NeuralNetworkParameters) = ps.params
+_get_params(ps::NeuralNetworkParameters) = params(ps)
 _get_params(ps::NamedTuple{(:params,), Tuple{NT}}) where {NT<:NamedTuple} = ps.params
 _get_params(ps::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}) = [_get_params(nt) for nt in ps]
 
diff --git a/src/symbolic_neuralnet/symbolic_neuralnet.jl b/src/symbolic_neuralnet/symbolic_neuralnet.jl
index da5a2c0..b3466e8 100644
--- a/src/symbolic_neuralnet/symbolic_neuralnet.jl
+++ b/src/symbolic_neuralnet/symbolic_neuralnet.jl
@@ -31,7 +31,7 @@ end
 
 function SymbolicNeuralNetwork(nn::NeuralNetwork)
     cache = Dict()
-    sparams = symbolize!(cache, nn.params, :W)
+    sparams = symbolize!(cache, params(nn), :W)
     @variables sinput[1:input_dimension(nn.model)]
 
     SymbolicNeuralNetwork(nn.architecture, nn.model, sparams, sinput)
@@ -54,6 +54,8 @@ function SymbolicNeuralNetwork(d::AbstractExplicitLayer)
     SymbolicNeuralNetwork(UnknownArchitecture(), d)
 end
 
+params(snn::AbstractSymbolicNeuralNetwork) = snn.params
+
 apply(snn::AbstractSymbolicNeuralNetwork, x, args...) = snn(x, args...)
 
 input_dimension(::AbstractExplicitLayer{M}) where M = M 
@@ -68,5 +70,5 @@ function Base.show(io::IO, snn::SymbolicNeuralNetwork)
     print(io, "\nModel = ")
     print(io, snn.model)
     print(io, "\nSymbolic Params = ")
-    print(io, snn.params)
+    print(io, params(snn))
 end
\ No newline at end of file
diff --git a/src/symbolic_neuralnet/symbolize.jl b/src/symbolic_neuralnet/symbolize.jl
index 9738059..abf6e46 100644
--- a/src/symbolic_neuralnet/symbolize.jl
+++ b/src/symbolic_neuralnet/symbolize.jl
@@ -95,15 +95,15 @@ And for neural network parameters:
 
 ```jldoctest
 using SymbolicNeuralNetworks: symbolize!
-using AbstractNeuralNetworks
+using AbstractNeuralNetworks: NeuralNetwork, params, Chain, Dense
 
 nn = NeuralNetwork(Chain(Dense(1, 2; use_bias = false), Dense(2, 1; use_bias = false)))
 cache = Dict()
-sym = symbolize!(cache, nn.params, :X) |> typeof
+sym = symbolize!(cache, params(nn), :X) |> typeof
 
 # output
 
-NeuralNetworkParameters{(:L1, :L2), Tuple{@NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}}, @NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}}}}
+AbstractNeuralNetworks.NeuralNetworkParameters{(:L1, :L2), Tuple{@NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}}, @NamedTuple{W::Symbolics.Arr{Symbolics.Num, 2}}}}
 ```
 
 # Implementation
@@ -128,5 +128,5 @@ function symbolize!(cache::Dict, nt::NamedTuple, var_name::Symbol; redundancy::B
 end
 
 function symbolize!(cache::Dict, nt::NeuralNetworkParameters, var_name::Symbol; redundancy::Bool = true)
-    NeuralNetworkParameters(symbolize!(cache, nt.params, var_name; redundancy = redundancy))
+    NeuralNetworkParameters(symbolize!(cache, params(nt), var_name; redundancy = redundancy))
 end
\ No newline at end of file
diff --git a/test/build_function/build_function.jl b/test/build_function/build_function.jl
index 5aef4c6..76200c9 100644
--- a/test/build_function/build_function.jl
+++ b/test/build_function/build_function.jl
@@ -1,6 +1,7 @@
 using SymbolicNeuralNetworks: _build_nn_function
 using SymbolicNeuralNetworks
 using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params
 using Test
 
 # this tests the function '_build_nn_function' (not 'build_nn_function')
@@ -8,11 +9,11 @@ function apply_build_function(input_dim::Integer=2, output_dim::Integer=1, num_e
     c = Chain(Dense(input_dim, output_dim, tanh))
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
-    eq = c(snn.input, snn.params)
-    built_function = _build_nn_function(eq, snn.params, snn.input)
+    eq = c(snn.input, params(snn))
+    built_function = _build_nn_function(eq, params(snn), snn.input)
     input = rand(input_dim, num_examples)
 
-    @test all(i -> (built_function(input, nn.params, i) ≈ c(input[:, i], nn.params)), 1:num_examples)
+    @test all(i -> (built_function(input, params(nn), i) ≈ c(input[:, i], params(nn))), 1:num_examples)
 end
 
 for input_dim ∈ (2, 3)
diff --git a/test/build_function/build_function_arrays.jl b/test/build_function/build_function_arrays.jl
index e9f0980..0da9bad 100644
--- a/test/build_function/build_function_arrays.jl
+++ b/test/build_function/build_function_arrays.jl
@@ -1,5 +1,5 @@
 using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork, function_valued_parameters
-using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
 using Test
 import Random
 Random.seed!(123)
@@ -8,13 +8,13 @@ function build_function_for_array_valued_equation(input_dim::Integer=2, output_d
     ch = Chain(Dense(input_dim, output_dim, tanh))
     nn = NeuralNetwork(ch)
     snn = SymbolicNeuralNetwork(nn)
-    eqs = [(a = ch(snn.input, snn.params), b = ch(snn.input, snn.params).^2), (c = ch(snn.input, snn.params).^3, )]
-    funcs = build_nn_function(eqs, snn.params, snn.input)
+    eqs = [(a = ch(snn.input, params(snn)), b = ch(snn.input, params(snn)).^2), (c = ch(snn.input, params(snn)).^3, )]
+    funcs = build_nn_function(eqs, params(snn), snn.input)
     input = Vector(1:input_dim)
-    a = ch(input, nn.params)
-    b = ch(input, nn.params).^2
-    c = ch(input, nn.params).^3
-    funcs_evaluated = funcs(input, nn.params)
+    a = ch(input, params(nn))
+    b = ch(input, params(nn)).^2
+    c = ch(input, params(nn)).^3
+    funcs_evaluated = funcs(input, params(nn))
     funcs_evaluated_as_vector = [funcs_evaluated[1].a, funcs_evaluated[1].b, funcs_evaluated[2].c]
     result_of_standard_computation = [a, b, c]
 
@@ -25,12 +25,12 @@ function build_function_for_named_tuple(input_dim::Integer=2, output_dim::Intege
     c = Chain(Dense(input_dim, output_dim, tanh))
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
-    eqs = (a = c(snn.input, snn.params), b = c(snn.input, snn.params).^2)
-    funcs = build_nn_function(eqs, snn.params, snn.input)
+    eqs = (a = c(snn.input, params(snn)), b = c(snn.input, params(snn)).^2)
+    funcs = build_nn_function(eqs, params(snn), snn.input)
     input = Vector(1:input_dim)
-    a = c(input, nn.params)
-    b = c(input, nn.params).^2
-    funcs_evaluated = funcs(input, nn.params)
+    a = c(input, params(nn))
+    b = c(input, params(nn)).^2
+    funcs_evaluated = funcs(input, params(nn))
 
     funcs_evaluated_as_vector = [funcs_evaluated.a, funcs_evaluated.b]
     result_of_standard_computation = [a, b]
@@ -42,13 +42,13 @@ function function_valued_parameters_for_named_tuple(input_dim::Integer=2, output
     c = Chain(Dense(input_dim, output_dim, tanh))
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
-    eqs = (a = c(snn.input, snn.params), b = c(snn.input, snn.params).^2)
-    funcs = function_valued_parameters(eqs, snn.params, snn.input)
+    eqs = (a = c(snn.input, params(snn)), b = c(snn.input, params(snn)).^2)
+    funcs = function_valued_parameters(eqs, params(snn), snn.input)
     input = Vector(1:input_dim)
-    a = c(input, nn.params)
-    b = c(input, nn.params).^2
+    a = c(input, params(nn))
+    b = c(input, params(nn)).^2
 
-    funcs_evaluated_as_vector = [funcs.a(input, nn.params), funcs.b(input, nn.params)]
+    funcs_evaluated_as_vector = [funcs.a(input, params(nn)), funcs.b(input, params(nn))]
     result_of_standard_computation = [a, b]
     
     @test funcs_evaluated_as_vector ≈ result_of_standard_computation
diff --git a/test/build_function/build_function_double_input.jl b/test/build_function/build_function_double_input.jl
index 01854d8..59228dd 100644
--- a/test/build_function/build_function_double_input.jl
+++ b/test/build_function/build_function_double_input.jl
@@ -2,6 +2,7 @@ using SymbolicNeuralNetworks: _build_nn_function
 using SymbolicNeuralNetworks
 using Symbolics: @variables
 using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params
 using Test
 
 # this tests the function '_build_nn_function' (not 'build_nn_function') for input and output
@@ -10,12 +11,12 @@ function apply_build_function(input_dim::Integer=2, output_dim::Integer=1, num_e
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
     @variables soutput[1:output_dim]
-    eq = (c(snn.input, snn.params) - soutput).^2
-    built_function = _build_nn_function(eq, snn.params, snn.input, soutput)
+    eq = (c(snn.input, params(snn)) - soutput).^2
+    built_function = _build_nn_function(eq, params(snn), snn.input, soutput)
     input = rand(input_dim, num_examples)
     output = rand(output_dim, num_examples)
 
-    @test all(i -> (built_function(input, output, nn.params, i) ≈ (c(input[:, i], nn.params) - output[:, i]).^2), 1:num_examples)
+    @test all(i -> (built_function(input, output, params(nn), i) ≈ (c(input[:, i], params(nn)) - output[:, i]).^2), 1:num_examples)
 end
 
 for input_dim ∈ (2, 3)
diff --git a/test/derivatives/pullback.jl b/test/derivatives/pullback.jl
index 9eee5c7..54e7591 100644
--- a/test/derivatives/pullback.jl
+++ b/test/derivatives/pullback.jl
@@ -1,6 +1,7 @@
 using SymbolicNeuralNetworks
 using SymbolicNeuralNetworks: _get_params, _get_contents
 using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params
 using Symbolics
 using GeometricMachineLearning: ZygotePullback
 using Test
@@ -20,12 +21,12 @@ function compare_symbolic_pullback_to_zygote_pullback(input_dim::Integer, output
     loss = FeedForwardLoss()
     spb = SymbolicPullback(snn, loss)
     input_output = (rand(input_dim, second_dim), rand(output_dim, second_dim))
-    loss_and_pullback = spb(nn.params, nn.model, input_output)
+    loss_and_pullback = spb(params(nn), nn.model, input_output)
     # note that we apply the second argument to another input `1`
     pb_values = loss_and_pullback[2](1)
 
     zpb = ZygotePullback(loss)
-    loss_and_pullback_zygote = zpb(nn.params, nn.model, input_output)
+    loss_and_pullback_zygote = zpb(params(nn), nn.model, input_output)
     pb_values_zygote = loss_and_pullback_zygote[2](1) |> _get_contents |> _get_params
 
     compare_values(pb_values, pb_values_zygote)
diff --git a/test/derivatives/symbolic_gradient.jl b/test/derivatives/symbolic_gradient.jl
index 028d8c5..0b3ae4d 100644
--- a/test/derivatives/symbolic_gradient.jl
+++ b/test/derivatives/symbolic_gradient.jl
@@ -2,7 +2,7 @@ using SymbolicNeuralNetworks
 using SymbolicNeuralNetworks: symbolic_differentials, symbolic_derivative, _build_nn_function
 using LinearAlgebra: norm
 using Symbolics, AbstractNeuralNetworks
-using AbstractNeuralNetworks: NeuralNetworkParameters
+using AbstractNeuralNetworks: NeuralNetworkParameters, params
 
 import Zygote
 import Random
@@ -16,16 +16,16 @@ function test_symbolic_gradient(input_dim::Integer = 3, output_dim::Integer = 1,
     c = Chain(Dense(input_dim, hidden_dim, tanh), Dense(hidden_dim, output_dim, tanh))
     nn = NeuralNetwork(c)
     snn = SymbolicNeuralNetwork(nn)
-    sout = norm(c(snn.input, snn.params)) ^ 2
-    sdparams = symbolic_differentials(snn.params)
+    sout = norm(c(snn.input, params(snn))) ^ 2
+    sdparams = symbolic_differentials(params(snn))
     _sgrad = symbolic_derivative(sout, sdparams)
     input = rand(T, input_dim, second_dim)
     for k in 1:second_dim
-        zgrad = Zygote.gradient(ps -> (norm(c(input[:, k], ps)) ^ 2), nn.params)[1].params
+        zgrad = Zygote.gradient(ps -> (norm(c(input[:, k], ps)) ^ 2), params(nn))[1].params
         for key1 in keys(_sgrad)
             for key2 in keys(_sgrad[key1])
-                executable_gradient = _build_nn_function(_sgrad[key1][key2], snn.params, snn.input)
-                sgrad = executable_gradient(input, nn.params, k)
+                executable_gradient = _build_nn_function(_sgrad[key1][key2], params(snn), snn.input)
+                sgrad = executable_gradient(input, params(nn), k)
                 @test sgrad ≈ zgrad[key1][key2]
             end
         end
@@ -40,12 +40,12 @@ function test_symbolic_gradient2(input_dim::Integer = 3, output_dim::Integer = 1
     c = Chain(Dense(input_dim, hidden_dim, tanh), Dense(hidden_dim, output_dim, tanh))
     nn = NeuralNetwork(c, T)
     snn = SymbolicNeuralNetwork(nn)
-    sout = norm(c(snn.input, snn.params)) ^ 2
+    sout = norm(c(snn.input, params(snn))) ^ 2
     input = rand(T, input_dim, second_dim, third_dim)
-    zgrad = Zygote.gradient(ps -> (norm(c(input, ps)) ^ 2), nn.params)[1].params
-    sdparams = symbolic_differentials(snn.params)
+    zgrad = Zygote.gradient(ps -> (norm(c(input, ps)) ^ 2), params(nn))[1].params
+    sdparams = symbolic_differentials(params(snn))
     _sgrad = symbolic_derivative(sout, sdparams)
-    sgrad = build_nn_function(_sgrad, sparams, sinput)(input, nn.params)
+    sgrad = build_nn_function(_sgrad, sparams, sinput)(input, params(nn))
     for key1 in keys(sgrad) for key2 in keys(sgrad[key1]) @test zgrad[key1][key2] ≈ sgrad[key1][key2] end end
 end
 

From 94e9da1898226ecfcaf080d0e278fef9817bba39 Mon Sep 17 00:00:00 2001
From: benedict-96 <benedikt.brantner@icloud.com>
Date: Tue, 17 Dec 2024 13:39:45 +0100
Subject: [PATCH 23/23] Resolved merge conflict.

---
 src/build_function/build_function_arrays.jl | 11 ++++++-----
 src/derivatives/pullback.jl                 |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/build_function/build_function_arrays.jl b/src/build_function/build_function_arrays.jl
index 49f2b10..6b25877 100644
--- a/src/build_function/build_function_arrays.jl
+++ b/src/build_function/build_function_arrays.jl
@@ -83,16 +83,17 @@ Return an executable function for each entry in `eqs`. This still has to be proc
 
 ```jldoctest
 using SymbolicNeuralNetworks: function_valued_parameters, SymbolicNeuralNetwork
-using AbstractNeuralNetworks: Chain, Dense,fffff, NeuralNetworkParameters, params
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
 import Random
 Random.seed!(123)
 
 c = Chain(Dense(2, 1, tanh))
-nn = SymbolicNeuralNetwork(c)
-eqs = (a = c(nn.input, params(nn)), b = c(nn.input, params(nn)).^2)
-funcs = function_valued_parameters(eqs, params(nn), nn.input)
+nn = NeuralNetwork(c)
+snn = SymbolicNeuralNetwork(nn)
+eqs = (a = c(snn.input, params(snn)), b = c(snn.input, params(snn)).^2)
+funcs = function_valued_parameters(eqs, params(snn), snn.input)
 input = [1., 2.]
-ps = initialparameters(c) |> NeuralNetworkParameters
+ps = params(nn)
 a = c(input, ps)
 b = c(input, ps).^2
 
diff --git a/src/derivatives/pullback.jl b/src/derivatives/pullback.jl
index 50f8a6a..803c34a 100644
--- a/src/derivatives/pullback.jl
+++ b/src/derivatives/pullback.jl
@@ -143,4 +143,4 @@ _get_contents(nt::Tuple{<:Union{NamedTuple, NeuralNetworkParameters}}) = nt[1]
 # (_pullback::SymbolicPullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps)
 function (_pullback::SymbolicPullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple
     _pullback.loss(model, ps, input_nt_output_nt...), _pullback.fun(input_nt_output_nt..., ps)
-end
\ No newline at end of file
+end