Merge b3bb9af into c768c0b

Author: prbzrg

src/utils.jl

@@ -6,7 +6,12 @@ function jacobian_batched(
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
+    res = Zygote.Buffer(
+        convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
+        size(xs, 1),
+        size(xs, 1),
+        size(xs, 2),
+    )
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[i, :, :] =
@@ -24,7 +29,12 @@ function jacobian_batched(
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
+    res = Zygote.Buffer(
+        convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
+        size(xs, 1),
+        size(xs, 1),
+        size(xs, 2),
+    )
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[:, i, :] = only(