more static removal

Author: ChrisRackauckas

test/fast_neural_ode.jl

@@ -5,7 +5,7 @@ datasize = 30
 tspan = (0.0f0,1.5f0)
 
 function trueODEfunc(du,u,p,t)
-    true_A = [-0.1 2.0; -2.0 -0.1]
+    true_A = Float32[-0.1 2.0; -2.0 -0.1]
     du .= ((u.^3)'true_A)'
 end
 t = range(tspan[1],tspan[2],length=datasize)
@@ -27,21 +27,6 @@ function fast_loss_n_ode(p)
     loss,pred
 end
 
-staticdudt2 = FastChain((x,p) -> x.^3,
-                        StaticDense(2,50,tanh),
-                        StaticDense(50,2))
-static_n_ode = NeuralODE(staticdudt2,tspan,Tsit5(),saveat=t)
-
-function static_predict_n_ode(p)
-  static_n_ode(u0,p)
-end
-
-function static_loss_n_ode(p)
-    pred = static_predict_n_ode(p)
-    loss = sum(abs2,ode_data .- pred)
-    loss,pred
-end
-
 dudt2 = Flux.Chain((x) -> x.^3,
                 Flux.Dense(2,50,tanh),
                 Flux.Dense(50,2))
@@ -60,12 +45,5 @@ end
 p = initial_params(fastdudt2)
 _p,re = Flux.destructure(dudt2)
 @test fastdudt2(ones(2),_p) ≈ dudt2(ones(2))
-@test staticdudt2(ones(2),_p) ≈ dudt2(ones(2))
 @test fast_loss_n_ode(p)[1] ≈ loss_n_ode(p)[1]
-@test static_loss_n_ode(p)[1] ≈ loss_n_ode(p)[1]
-@test Zygote.gradient((p)->fast_loss_n_ode(p)[1], p)[1] ≈ Zygote.gradient((p)->loss_n_ode(p)[1], p)[1] rtol=4e-3
-@test Zygote.gradient((p)->static_loss_n_ode(p)[1], p)[1] ≈ Zygote.gradient((p)->loss_n_ode(p)[1], p)[1] rtol=4e-3
-
-# @btime Zygote.gradient((p)->static_loss_n_ode(p)[1], p)
-# @btime Zygote.gradient((p)->fast_loss_n_ode(p)[1], p)
-# @btime Zygote.gradient((p)->loss_n_ode(p)[1], p)
+@test Zygote.gradient((p)->fast_loss_n_ode(p)[1], p)[1] ≈ Zygote.gradient((p)->loss_n_ode(p)[1], p)[1] rtol=4e-3