Merge branch 'main' into output-to-f

Author: benedict-96

Project.toml

@@ -28,12 +28,12 @@ julia = "1.10"
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GeometricMachineLearning = "194d25b2-d3f5-49f0-af24-c124f4aa80cc"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 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", "GeometricMachineLearning"]

src/build_function/build_function.jl

@@ -69,31 +69,16 @@ This first calls `Symbolics.build_function` with the keyword argument `expressio
 
 See the docstrings for those functions for details on how the code is modified. 
 """
-function _build_nn_function(eq::EqT, params::NeuralNetworkParameters, sinput::Symbolics.Arr)
+function _build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr)
     sc_eq = Symbolics.scalarize(eq)
-    code = build_function(sc_eq, sinput, values(params)...; expression = Val{true}) |> _reduce_code
+    code = build_function(sc_eq, sinput, values(sparams)...; expression = Val{true}) |> _reduce
     rewritten_code = fix_map_reduce(modify_input_arguments(rewrite_arguments(fix_create_array(code))))
     parallelized_code = make_kernel(rewritten_code)
     @RuntimeGeneratedFunction(parallelized_code)
 end
 
-"""
-    _reduce_code(code)
-
-Reduce the code.
-
-For some reason `Symbolics.build_function` sometimes returns a tuple and sometimes it doesn't.
-
-This function takes care of this. 
-If `build_function` returns a tuple `reduce_code` checks which of the expressions is in-place and then returns the other (not in-place) expression.
-"""
-function _reduce_code(code::Expr)
-    code
-end
-
-function _reduce_code(code::Tuple{Expr, Expr})
-    contains(string(code[1]), "ˍ₋out") ? code[2] : code[1]
-end
+_reduce(a) = a
+_reduce(a::Tuple) = a[1]
 
 """
     rewrite_arguments(s)

src/build_function/build_function_double_input.jl

@@ -42,7 +42,7 @@ See the docstrings for those functions for details on how the code is modified.
 """
 function _build_nn_function(eq::EqT, params::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr)
     sc_eq = Symbolics.scalarize(eq)
-    code = build_function(sc_eq, sinput, soutput, values(params)...; expression = Val{true}) |> _reduce_code
+    code = build_function(sc_eq, sinput, soutput, values(params)...; expression = Val{true}) |> _reduce
     rewritten_code = fix_map_reduce(modify_input_arguments2(rewrite_arguments2(fix_create_array(code))))
     parallelized_code = make_kernel2(rewritten_code)
     @RuntimeGeneratedFunction(parallelized_code)

src/derivatives/derivative.jl

@@ -1,7 +1,7 @@
 """
     Derivative
 """
-abstract type Derivative{ST, FT, SDT} end
+abstract type Derivative{OT, SDT, ST <: AbstractSymbolicNeuralNetwork} end
 
 derivative(::DT) where {DT <: Derivative} = error("No method of function `derivative` defined for type $(DT).")
 

src/derivatives/gradient.jl

@@ -32,10 +32,10 @@ nn = SymbolicNeuralNetwork(c)
 
 Internally the constructors are using [`symbolic_pullback`](@ref).
 """
-struct Gradient{ST, FT, SDT} <: Derivative{ST, FT, SDT} 
-    nn::ST
-    f::FT
+struct Gradient{OT, SDT, ST} <: Derivative{OT, SDT, ST} 
+    output::OT
     ∇::SDT
+    nn::ST
 end
 
 """
@@ -65,7 +65,7 @@ derivative(g::Gradient) = g.∇
 
 function Gradient(output::EqT, nn::SymbolicNeuralNetwork)
     typeof(output) <: AbstractArray ? nothing : (@warn "You should only use `Gradient` together with array expressions! Maybe you wanted to use `SymbolicPullback`.")
-    Gradient(nn, output, symbolic_pullback(output, nn))
+    Gradient(output, symbolic_pullback(output, nn), nn)
 end
 
 function Gradient(nn::SymbolicNeuralNetwork)
@@ -86,14 +86,13 @@ 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, params(nn))
 spb = symbolic_pullback(output, nn)
 
-spb[1].L1.b |> latexify
+spb[1].L1.b
 ```
 """
 function symbolic_pullback(f::EqT, nn::AbstractSymbolicNeuralNetwork)::Union{AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, Union{NamedTuple, NeuralNetworkParameters}}

src/derivatives/jacobian.jl

@@ -74,28 +74,28 @@ jacobian1(input, ps) ≈ [analytic_jacobian(i, j) for j ∈ 1:output_dim, i ∈
 true
 ```
 """
-struct Jacobian{ST, FT, SDT} <: Derivative{ST, FT, SDT} 
-    nn::ST
-    f::FT
+struct Jacobian{OT, SDT, ST} <: Derivative{OT, SDT, ST} 
+    output::OT
     □::SDT
+    nn::ST
 end
 
 derivative(j::Jacobian) = j.□
 
-function Jacobian(nn::AbstractSymbolicNeuralNetwork)
-    
-    # Evaluation of the symbolic output
-    soutput = nn.model(nn.input, params(nn))
-
-    Jacobian(soutput, nn)
-end
-
 function Jacobian(f::EqT, nn::AbstractSymbolicNeuralNetwork)
     # make differential 
     Dx = symbolic_differentials(nn.input)
 
     # Evaluation of gradient
     s∇f = hcat([expand_derivatives.(Symbolics.scalarize(dx(f))) for dx in Dx]...)
 
-    Jacobian(nn, f, s∇f)
+    Jacobian(f, s∇f, nn)
+end
+
+function Jacobian(nn::AbstractSymbolicNeuralNetwork)
+    
+    # Evaluation of the symbolic output
+    soutput = nn.model(nn.input, params(nn))
+
+    Jacobian(soutput, nn)
 end

test/derivatives/jacobian.jl

@@ -32,7 +32,7 @@ function test_jacobian(n::Integer, T = Float32)
     @test build_nn_function(derivative(g), nn)(input, _params) ≈ ForwardDiff.jacobian(input -> c(input, _params), input)
 end
 
-for n ∈ 1:10
+for n ∈ 10:1
     for T ∈ (Float32, Float64)
         test_jacobian(n, T)
     end