Add JET tests

.github/workflows/JET.yml

@@ -0,0 +1,60 @@
+name: JET
+on:
+  push:
+    branches:
+      - master
+    tags: ['*']
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  test:
+    if: "!contains(github.event.head_commit.message, 'skip ci')"
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.12'
+        os:
+          - ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+      - run: julia -e 'using InteractiveUtils; versioninfo(verbose=true)'
+      - uses: julia-actions/cache@v2
+      - name: Install and run JET.jl
+        shell: julia --project="." --color=yes {0}
+        run: |
+          using Pkg
+          Pkg.add("JET")
+          Pkg.instantiate()
+          using JET
+          using RecursiveArrayTools
+
+          # Get all reports first
+          result = JET.report_package(RecursiveArrayTools; target_modules = (RecursiveArrayTools,))
+          reports = JET.get_reports(result)
+
+          # Filter out similar_type inference errors from StaticArraysCore
+          filtered_reports = filter(reports) do report
+              s = string(report)
+              !(occursin("similar_type", s) && occursin("StaticArraysCore", s))
+          end
+
+          # Check if there are any non-filtered errors
+          if !isempty(filtered_reports)
+              @error "JET found errors" filtered_reports
+              exit(1)
+          else
+              @info "All JET errors are filtered (similar_type related)"
+          end

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

@@ -924,6 +924,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)
@@ -1200,6 +1202,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