Merge pull request #301 from avik-pal/patch-1

test: new enzyme version

Author: ChrisRackauckas

Project.toml

@@ -4,7 +4,7 @@ authors = [
     "Pankaj Mishra <pankajmishra1511@gmail.com>",
     "Chris Rackauckas <contact@chrisrackauckas.com>",
 ]
-version = "2.20.0"
+version = "2.21.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -48,7 +48,7 @@ Adapt = "3, 4"
 ArrayInterface = "7.4.2"
 Compat = "4"
 DataStructures = "0.18"
-Enzyme = "0.11,0.12"
+Enzyme = "0.13"
 FiniteDiff = "2.8.1"
 ForwardDiff = "0.10"
 Graphs = "1"

ext/SparseDiffToolsEnzymeExt.jl

@@ -27,7 +27,8 @@ function __gradient!(::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, f!, fx, x, c
 end
 
 function __jacobian!(J::AbstractMatrix, ::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, f, x)
-    J .= jacobian(Reverse, f, x, Val(size(J, 1)))
+    J .= only(jacobian(
+        Reverse, a -> vec(f(reshape(a, size(x)))), vec(x); n_outs = Val(size(J, 1))))
     return J
 end