Preserve the dimensions of ReverseDiff TrackedArray

JuliaArrays · Mar 7, 2024 · fc353f8 · fc353f8
1 parent bbf06e2
commit fc353f8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "7.8.0"
+version = "7.8.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/ext/ArrayInterfaceReverseDiffExt.jl b/ext/ArrayInterfaceReverseDiffExt.jl
@@ -14,9 +14,9 @@ ArrayInterface.can_setindex(::Type{<:ReverseDiff.TrackedArray}) = false
 ArrayInterface.fast_scalar_indexing(::Type{<:ReverseDiff.TrackedArray}) = false
 function ArrayInterface.aos_to_soa(x::AbstractArray{<:ReverseDiff.TrackedReal,N}) where {N}
     if length(x) > 1
-        reduce(vcat,x)
+        return reshape(reduce(vcat,x), size(x))
     else
-        reduce(vcat,[x[1],x[1]])[1:1]
+        return reduce(vcat,[x[1],x[1]])[1:1]
     end
 end