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

test: new enzyme version
JuliaDiff · Sep 26, 2024 · 0c6f216 · 0c6f216
2 parents 47bb1e2 + 155d05e
commit 0c6f216
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -4,7 +4,7 @@ authors = [
     "Pankaj Mishra <[email protected]>",
     "Chris Rackauckas <[email protected]>",
 ]
-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"

diff --git a/ext/SparseDiffToolsEnzymeExt.jl b/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