Merge pull request #25 from JuliaGNI/flip-arguments-in-jacobian

Flip keys in jacobian

Author: michakraus

Project.toml

@@ -3,16 +3,16 @@ uuid = "aed23131-dcd0-47ca-8090-d21e605652e3"
 authors = ["Michael Kraus"]
 version = "0.2.0"
 
-[weakdeps]
-SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
-
 [deps]
 AbstractNeuralNetworks = "60874f82-5ada-4c70-bd1c-fa6be7711c8a"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
+[weakdeps]
+SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
+
 [compat]
 AbstractNeuralNetworks = "0.3, 0.4, 0.5, 0.6"
 Documenter = "1.8.0"
@@ -28,12 +28,12 @@ Zygote = "0.6.73"
 [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/derivatives/derivative.jl

@@ -1,7 +1,7 @@
 """
     Derivative
 """
-abstract type Derivative{ST, OT, 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

@@ -28,10 +28,10 @@ nn = SymbolicNeuralNetwork(c)
 
 Internally the constructors are using [`symbolic_pullback`](@ref).
 """
-struct Gradient{ST, OT, SDT} <: Derivative{ST, OT, SDT} 
-    nn::ST
+struct Gradient{OT, SDT, ST} <: Derivative{OT, SDT, ST} 
     output::OT
     ∇::SDT
+    nn::ST
 end
 
 """
@@ -59,7 +59,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)
@@ -87,7 +87,7 @@ 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(soutput::EqT, nn::AbstractSymbolicNeuralNetwork)::Union{AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, Union{NamedTuple, NeuralNetworkParameters}}

src/derivatives/jacobian.jl

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