diff --git a/Project.toml b/Project.toml
index 4b05d7f..92e83be 100644
--- a/Project.toml
+++ b/Project.toml
@@ -14,9 +14,9 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 AbstractNeuralNetworks = "0.3, 0.4, 0.5"
 Documenter = "1.8.0"
 ForwardDiff = "0.10.38"
+GeometricMachineLearning = "0.3.7"
 Latexify = "0.16.5"
 RuntimeGeneratedFunctions = "0.5"
-SafeTestsets = "0.1"
 Symbolics = "5, 6"
 Zygote = "0.6.73"
 julia = "1.6"
@@ -29,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/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/hamiltonian_neural_network.md b/docs/src/hamiltonian_neural_network.md
index 17c046b..71f8af4 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)
diff --git a/docs/src/symbolic_neural_networks.md b/docs/src/symbolic_neural_networks.md
index 9de1b33..32a3de1 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
+using AbstractNeuralNetworks: Chain, Dense, 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/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl
index c7d8294..49ea427 100644
--- a/src/SymbolicNeuralNetworks.jl
+++ b/src/SymbolicNeuralNetworks.jl
@@ -15,42 +15,26 @@ module SymbolicNeuralNetworks
     
     RuntimeGeneratedFunctions.init(@__MODULE__)
 
-    include("equation_types.jl")
+    include("custom_definitions_and_extensions/equation_types.jl")
 
-    export symbolize
-    include("utils/symbolize.jl")
+    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.jl")
-
-    export evaluate_equations
-    include("symbolic_neuralnet.jl")
-
-    export symbolic_hamiltonian
-    include("hamiltonian.jl")
+    include("symbolic_neuralnet/symbolic_neuralnet.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_function_double_input.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")
     include("derivatives/gradient.jl")
 
-    include("custom_equation.jl")
-
-    include("utils/latexraw.jl")
+    include("custom_definitions_and_extensions/latexraw.jl")
 end
diff --git a/src/utils/build_function.jl b/src/build_function/build_function.jl
similarity index 93%
rename from src/utils/build_function.jl
rename to src/build_function/build_function.jl
index bcefd91..4943989 100644
--- a/src/utils/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)
@@ -39,25 +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, symbolicparameters
-using Symbolics
-using AbstractNeuralNetworks
+using SymbolicNeuralNetworks: _build_nn_function, SymbolicNeuralNetwork
+using AbstractNeuralNetworks: params, Chain, Dense, NeuralNetwork
+import Random
+Random.seed!(123)
 
 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 = NeuralNetwork(c).params
-input = rand(2, 2)
-
-(built_function(input, ps, 1), built_function(input, ps, 2)) .≈ (c(input[:, 1], ps), c(input[:, 2], ps))
+nn = NeuralNetwork(c)
+snn = SymbolicNeuralNetwork(nn)
+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
 
-(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/src/utils/build_function_arrays.jl b/src/build_function/build_function_arrays.jl
similarity index 87%
rename from src/utils/build_function_arrays.jl
rename to src/build_function/build_function_arrays.jl
index 265c442..6b25877 100644
--- a/src/utils/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, NeuralNetwork
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
 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, 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.]
-ps = NeuralNetwork(ch).params
-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, params(nn))
 
 # 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...)
@@ -47,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, NeuralNetwork
+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, 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, params(snn)), b = c(snn.input, params(snn)).^2)
+funcs = build_nn_function(eqs, params(snn), snn.input)
 input = [1., 2.]
-ps = NeuralNetwork(c).params
-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, params(nn))
 
 # output
 
-(true, true)
+(a = [-0.9999386280616135], b = [0.9998772598897417])
 ```
 
 # Implementation
@@ -90,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, NeuralNetwork
+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, nn.params), b = c(nn.input, nn.params).^2)
-funcs = function_valued_parameters(eqs, nn.params, 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 = NeuralNetwork(c).params
+ps = params(nn)
 a = c(input, ps)
 b = c(input, ps).^2
 
diff --git a/src/utils/build_function2.jl b/src/build_function/build_function_double_input.jl
similarity index 98%
rename from src/utils/build_function2.jl
rename to src/build_function/build_function_double_input.jl
index 77e1853..ece85c7 100644
--- a/src/utils/build_function2.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/chain.jl b/src/chain.jl
deleted file mode 100644
index 419f039..0000000
--- a/src/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/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/derivatives/gradient.jl b/src/derivatives/gradient.jl
index b075201..8790a80 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}
@@ -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
@@ -105,7 +106,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}
@@ -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 1cb0978..c93bb0e 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/pullback.jl b/src/derivatives/pullback.jl
similarity index 76%
rename from src/pullback.jl
rename to src/derivatives/pullback.jl
index 09ed479..803c34a 100644
--- a/src/pullback.jl
+++ b/src/derivatives/pullback.jl
@@ -8,15 +8,17 @@
 ```jldoctest
 using SymbolicNeuralNetworks
 using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params
 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 = NeuralNetwork(c).params
-pv_values = pb(ps, nn.model, (rand(2), rand(1)))[2](1) |> typeof
+pb = SymbolicPullback(snn, loss)
+ps = params(nn)
+typeof(pb(ps, nn.model, (rand(2), rand(1)))[2](1))
 
 # output
 
@@ -41,25 +43,26 @@ 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)
 
 c = Chain(Dense(2, 1, tanh))
-nn = SymbolicNeuralNetwork(c)
+nn = NeuralNetwork(c)
+snn = SymbolicNeuralNetwork(nn)
 loss = FeedForwardLoss()
-pb = SymbolicPullback(nn, loss)
-ps = NeuralNetwork(c).params
+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(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, 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)
+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))
 
-pv_values == (pv_values2 |> SymbolicNeuralNetworks._get_params |> SymbolicNeuralNetworks._get_contents)
+pb_values == (pb_values2 |> SymbolicNeuralNetworks._get_contents |> SymbolicNeuralNetworks._get_params)
 
 # output
 
@@ -85,15 +88,11 @@ 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)
+    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
@@ -109,7 +108,8 @@ 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]
 
 """
@@ -130,16 +130,17 @@ _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::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
     _pullback.loss(model, ps, input_nt_output_nt...), _pullback.fun(input_nt_output_nt..., ps)
-end
\ No newline at end of file
+end
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
diff --git a/src/layers/abstract.jl b/src/layers/abstract.jl
deleted file mode 100644
index 05e34db..0000000
--- a/src/layers/abstract.jl
+++ /dev/null
@@ -1,20 +0,0 @@
-"""
-    symbolicparameters(model)
-
-Obtain the symbolic parameters of a neural network model.
-
-# Examples
-
-```jldoctest
-using SymbolicNeuralNetworks
-using AbstractNeuralNetworks
-
-d = Dense(4, 5, tanh)
-symbolicparameters(d) |> 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.jl b/src/symbolic_neuralnet/symbolic_neuralnet.jl
similarity index 74%
rename from src/symbolic_neuralnet.jl
rename to src/symbolic_neuralnet/symbolic_neuralnet.jl
index ae53736..b3466e8 100644
--- a/src/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, 
@@ -27,12 +29,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, params(nn), :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)
@@ -47,10 +54,12 @@ 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 
-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])
 
@@ -61,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
new file mode 100644
index 0000000..abf6e46
--- /dev/null
+++ b/src/symbolic_neuralnet/symbolize.jl
@@ -0,0 +1,132 @@
+"""
+    symboliccounter!(cache, arg; redundancy)
+
+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(cache) ? cache[arg] += 1 : cache[arg] = 1
+        nam = string(arg) * "_" * string(cache[arg])
+        Symbol(nam)
+    else
+        arg
+    end
+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)
+
+# 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: NeuralNetwork, params, Chain, Dense
+
+nn = NeuralNetwork(Chain(Dense(1, 2; use_bias = false), Dense(2, 1; use_bias = false)))
+cache = Dict()
+sym = symbolize!(cache, params(nn), :X) |> typeof
+
+# output
+
+AbstractNeuralNetworks.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!(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!(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
+
+function symbolize!(cache::Dict, nt::NeuralNetworkParameters, var_name::Symbol; redundancy::Bool = true)
+    NeuralNetworkParameters(symbolize!(cache, params(nt), var_name; redundancy = redundancy))
+end
\ No newline at end of file
diff --git a/src/utils/symbolize.jl b/src/utils/symbolize.jl
deleted file mode 100644
index 45b7dae..0000000
--- a/src/utils/symbolize.jl
+++ /dev/null
@@ -1,48 +0,0 @@
-#=
-    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. 
-=#
-
-function SymbolicName(arg, storage; redundancy = true)
-    if redundancy
-        arg ∈ keys(storage) ? storage[arg] += 1 : storage[arg] = 1
-        nam = string(arg)*"_"*string(storage[arg])
-        return Symbol(nam)
-    else
-        nam = string(arg)
-        return Symbol(nam)
-    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
-
-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)
-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
-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
-end
-
-#symbolize(nn::NeuralNetwork) = symbolize(nn.params, missing, Dict())[1]
\ No newline at end of file
diff --git a/test/build_function/build_function.jl b/test/build_function/build_function.jl
new file mode 100644
index 0000000..76200c9
--- /dev/null
+++ b/test/build_function/build_function.jl
@@ -0,0 +1,25 @@
+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')
+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, params(snn))
+    built_function = _build_nn_function(eq, params(snn), snn.input)
+    input = rand(input_dim, 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)
+    for output_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_arrays.jl b/test/build_function/build_function_arrays.jl
new file mode 100644
index 0000000..0da9bad
--- /dev/null
+++ b/test/build_function/build_function_arrays.jl
@@ -0,0 +1,64 @@
+using SymbolicNeuralNetworks: build_nn_function, SymbolicNeuralNetwork, function_valued_parameters
+using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
+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, 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, 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]
+
+    @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, 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, 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]
+
+    @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, 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, params(nn))
+    b = c(input, params(nn)).^2
+
+    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
+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/build_function/build_function_double_input.jl b/test/build_function/build_function_double_input.jl
new file mode 100644
index 0000000..59228dd
--- /dev/null
+++ b/test/build_function/build_function_double_input.jl
@@ -0,0 +1,28 @@
+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
+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)
+    @variables soutput[1:output_dim]
+    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, params(nn), i) ≈ (c(input[:, i], params(nn)) - output[:, i]).^2), 1:num_examples)
+end
+
+for input_dim ∈ (2, 3)
+    for output_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/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/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/derivatives/pullback.jl b/test/derivatives/pullback.jl
new file mode 100644
index 0000000..54e7591
--- /dev/null
+++ b/test/derivatives/pullback.jl
@@ -0,0 +1,39 @@
+using SymbolicNeuralNetworks
+using SymbolicNeuralNetworks: _get_params, _get_contents
+using AbstractNeuralNetworks
+using AbstractNeuralNetworks: params
+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(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(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)
+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/symbolic_gradient.jl b/test/derivatives/symbolic_gradient.jl
similarity index 72%
rename from test/symbolic_gradient.jl
rename to test/derivatives/symbolic_gradient.jl
index 0159648..0b3ae4d 100644
--- a/test/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
@@ -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 = NeuralNetwork(c, T).params
-    @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, 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), ps)[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], sparams, sinput)
-                sgrad = executable_gradient(input, ps, 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
@@ -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 = NeuralNetwork(c, T).params
-    @variables sinput[1:input_dim]
-    sout = norm(c(sinput, sparams)) ^ 2
+    nn = NeuralNetwork(c, T)
+    snn = SymbolicNeuralNetwork(nn)
+    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), ps)[1].params
-    sdparams = symbolic_differentials(sparams)
+    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, ps)
+    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
 
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 8daf81c..200da87 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("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 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 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
+@safetestset "Compare Zygote Pullback with Symbolic Pullback                                         " begin include("derivatives/pullback.jl") end
\ No newline at end of file
diff --git a/test/symbolic_neuralnet/symbolize.jl b/test/symbolic_neuralnet/symbolize.jl
new file mode 100644
index 0000000..36eac72
--- /dev/null
+++ b/test/symbolic_neuralnet/symbolize.jl
@@ -0,0 +1,27 @@
+using AbstractNeuralNetworks: NeuralNetworkParameters
+using SymbolicNeuralNetworks
+using SymbolicNeuralNetworks: symbolize!
+using Symbolics
+using Test
+
+# a tuple of `NamedTuple`s and `Tuple`s.
+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)
+        ) )
+
+cache = Dict()
+sparams = symbolize!(cache, params, :W)
+
+@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 = 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
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_params.jl b/test/test_params.jl
deleted file mode 100644
index d3bc141..0000000
--- a/test/test_params.jl
+++ /dev/null
@@ -1,22 +0,0 @@
-using SymbolicNeuralNetworks
-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)
-            )
-
-sparams = symbolize(params)[1]
-
-@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
-
-verified_sparams = (    (W = W_1, b = b_1), 
-                        (W = (W_2 , W_3), c = c_1), 
-                        M_1, 
-                        (X_1, X_2)
-                        )
-
-@test sparams === verified_sparams
\ 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)