fix some of the issues found by JET.jl

Author: JoshuaLampert

src/named_array_partition.jl

@@ -1,6 +1,6 @@
 """
     NamedArrayPartition(; kwargs...)
-    NamedArrayPartition(x::NamedTuple) 
+    NamedArrayPartition(x::NamedTuple)
 
 Similar to an `ArrayPartition` but the individual arrays can be accessed via the
 constructor-specified names. However, unlike `ArrayPartition`, each individual array
@@ -22,7 +22,7 @@ function NamedArrayPartition(x::NamedTuple)
     return NamedArrayPartition(ArrayPartition{T, S}(values(x)), names_to_indices)
 end
 
-# Note: overloading `getproperty` means we cannot access `NamedArrayPartition` 
+# Note: overloading `getproperty` means we cannot access `NamedArrayPartition`
 # fields except through `getfield` and accessor functions.
 ArrayPartition(x::NamedArrayPartition) = getfield(x, :array_partition)
 
@@ -53,7 +53,7 @@ end
 function Base.similar(
         A::NamedArrayPartition, ::Type{T}, ::Type{S}, R::DataType...) where {T, S}
     NamedArrayPartition(
-        similar(getfield(A, :array_partition), T, S, R), getfield(A, :names_to_indices))
+    similar(getfield(A, :array_partition), T, S, R...), getfield(A, :names_to_indices))
 end
 
 Base.Array(x::NamedArrayPartition) = Array(ArrayPartition(x))
@@ -68,7 +68,7 @@ function Base.getproperty(x::NamedArrayPartition, s::Symbol)
     getindex(ArrayPartition(x).x, getproperty(getfield(x, :names_to_indices), s))
 end
 
-# this enables x.s = some_array. 
+# this enables x.s = some_array.
 @inline function Base.setproperty!(x::NamedArrayPartition, s::Symbol, v)
     index = getproperty(getfield(x, :names_to_indices), s)
     ArrayPartition(x).x[index] .= v

src/utils.jl

@@ -119,7 +119,8 @@ function recursivefill!(b::AbstractArray{T, N},
         a::T2) where {T <: StaticArraysCore.SArray,
         T2 <: Union{Number, Bool}, N}
     @inbounds for i in eachindex(b)
-        b[i] = fill(a, typeof(b[i]))
+        # Preserve static array shape while replacing all entries with the scalar
+        b[i] = map(_ -> a, b[i])
     end
 end
 
@@ -128,7 +129,8 @@ function recursivefill!(bs::AbstractVectorOfArray{T, N},
         T2 <: Union{Number, Bool}, N}
     @inbounds for b in bs, i in eachindex(b)
 
-        b[i] = fill(a, typeof(b[i]))
+        # Preserve static array shape while replacing all entries with the scalar
+        b[i] = map(_ -> a, b[i])
     end
 end
 

src/vector_of_array.jl

@@ -669,6 +669,8 @@ function Base.view(A::AbstractVectorOfArray{T, N, <:AbstractVector{T}},
         J = map(i -> Base.unalias(A, i), to_indices(A, I))
     elseif length(I) == 2 && (I[1] == Colon() || I[1] == 1)
         J = map(i -> Base.unalias(A, i), to_indices(A, Base.tail(I)))
+    else
+        J = map(i -> Base.unalias(A, i), to_indices(A, I))
     end
     @boundscheck checkbounds(A, J...)
     SubArray(A, J)
@@ -939,6 +941,7 @@ end
 
 struct VectorOfArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end # N is only used when voa sees other abstract arrays
 VectorOfArrayStyle{N}(::Val{N}) where {N} = VectorOfArrayStyle{N}()
+VectorOfArrayStyle(::Val{N}) where {N} = VectorOfArrayStyle{N}()
 
 # The order is important here. We want to override Base.Broadcast.DefaultArrayStyle to return another Base.Broadcast.DefaultArrayStyle.
 Broadcast.BroadcastStyle(a::VectorOfArrayStyle, ::Base.Broadcast.DefaultArrayStyle{0}) = a