Merge pull request #27 from JuliaGNI/remove-reduce-code

Remove reduce code

Author: michakraus

Project.toml

@@ -24,6 +24,7 @@ SafeTestsets = "0.1"
 SymbolicUtils = "<3.8.0"
 Symbolics = "5, 6"
 Zygote = "0.6.73"
+julia = "1.10"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

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/gradient.jl

@@ -80,7 +80,6 @@ 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)

test/derivatives/jacobian.jl

@@ -28,11 +28,13 @@ function test_jacobian(n::Integer, T = Float32)
 
     _params = params(NeuralNetwork(c, T))
     input = rand(T, n)
-    @test build_nn_function(g.output, nn)(input, _params) ≈ c(input, _params)
-    @test build_nn_function(derivative(g), nn)(input, _params) ≈ ForwardDiff.jacobian(input -> c(input, _params), input)
+    f = build_nn_function(g.output, nn)
+    ∇f = build_nn_function(derivative(g), nn)
+    @test f(input, params) ≈ c(input, params)
+    @test ∇f(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