Merge pull request #125 from ChrisRackauckas-Claude/compat-functors-0.5

Bump Functors compat to 0.5 and fix Zygote 0.7 compatibility

Author: ChrisRackauckas

Project.toml

@@ -29,7 +29,7 @@ DiffEqBase = "6.151"
 Distributions = "v0.25.107"
 DocStringExtensions = "0.9.3"
 Flux = "0.14.16, 0.15, 0.16"
-Functors = "0.4.11"
+Functors = "0.4.11, 0.5"
 LinearAlgebra = "1.10"
 Random = "1.10"
 Reexport = "1.2.2"

src/DeepBSDE.jl

@@ -165,7 +165,7 @@ function DiffEqBase.solve(
         )
         return map(sol) do _sol
             predict_ans = Array(_sol)
-            (predict_ans[1:(end - 1), end], predict_ans[end, end])
+            predict_ans[:, end]
         end
     end
 
@@ -176,7 +176,8 @@ function DiffEqBase.solve(
     end
 
     function loss_n_sde()
-        return mean(sum(abs2, g(X) - u) for (X, u) in predict_n_sde())
+        preds = predict_n_sde()
+        return mean(sum(abs2, g(pred[1:(end - 1)]) - pred[end]) for pred in preds)
     end
 
     iters = eltype(x0)[]

src/NNStopping.jl

@@ -47,7 +47,7 @@ struct NNStoppingModelArray{M}
     ms::M
 end
 
-Flux.@functor NNStoppingModelArray
+@functor NNStoppingModelArray
 
 function (model::NNStoppingModelArray)(X, G)
     XG = cat(X, reshape(G, 1, size(G)...), dims = 1)