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

fix: aos_to_soa for all singleton dims

Author: ChrisRackauckas

ext/ArrayInterfaceReverseDiffExt.jl

@@ -8,11 +8,8 @@ ArrayInterface.ismutable(T::Type{<:ReverseDiff.TrackedReal}) = false
 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
-        return reshape(reduce(vcat, x), size(x))
-    else
-        return reduce(vcat,[x[1], x[1]])[1:1]
-    end
+    y = length(x) > 1 ? reduce(vcat, x) : reduce(vcat, [x[1], x[1]])[1:1]
+    return reshape(y, size(x))
 end
 
 function ArrayInterface.restructure(x::Array, y::ReverseDiff.TrackedArray)

test/ad.jl

@@ -1,6 +1,9 @@
 using ArrayInterface, ReverseDiff, Tracker, Test
 x = ReverseDiff.track([4.0])
 @test ArrayInterface.aos_to_soa(x) isa ReverseDiff.TrackedArray
+x = reshape([ReverseDiff.track(rand(1, 1, 1))[1]], 1, 1, 1)
+@test ArrayInterface.aos_to_soa(x) isa ReverseDiff.TrackedArray
+@test ndims(ArrayInterface.aos_to_soa(x)) == 3
 x = reduce(vcat, ReverseDiff.track([4.0,4.0]))
 @test ArrayInterface.aos_to_soa(x) isa ReverseDiff.TrackedArray
 x = [ReverseDiff.track([4.0])[1]]