Preserve the dimensions of ReverseDiff TrackedArray

Author: avik-pal

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"

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