fix: aos_to_soa for all singleton dims

JuliaArrays · Nov 18, 2024 · d6f380a · d6f380a
1 parent 2643104
commit d6f380a
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/ext/ArrayInterfaceReverseDiffExt.jl b/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)