Merge pull request #510 from ChrisRackauckas-Claude/fix-raggedend-broadcasting

Make RaggedEnd and RaggedRange broadcast as scalars

Author: ChrisRackauckas

src/vector_of_array.jl

@@ -446,6 +446,10 @@ Base.:+(re::RaggedEnd, n::Integer) = RaggedEnd(re.dim, re.offset + Int(n))
 Base.:-(re::RaggedEnd, n::Integer) = RaggedEnd(re.dim, re.offset - Int(n))
 Base.:+(n::Integer, re::RaggedEnd) = re + n
 
+# Make RaggedEnd and RaggedRange broadcast as scalars to avoid
+# issues with collect/length in broadcasting contexts (e.g., SymbolicIndexingInterface)
+Base.broadcastable(x::RaggedEnd) = Ref(x)
+
 struct RaggedRange
     dim::Int
     start::Int
@@ -460,6 +464,7 @@ end
 function Base.:(:)(start::Integer, step::Integer, stop::RaggedEnd)
     RaggedRange(stop.dim, Int(start), Int(step), stop.offset)
 end
+Base.broadcastable(x::RaggedRange) = Ref(x)
 
 @inline function _is_ragged_dim(VA::AbstractVectorOfArray, d::Integer)
     length(VA.u) <= 1 && return false

test/basic_indexing.jl

@@ -196,6 +196,20 @@ ragged2 = VectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]])
 @test ragged2[1:(end - 1), 2] == [5.0]
 @test ragged2[1:(end - 1), 3] == [7.0, 8.0]
 
+# Test that RaggedEnd and RaggedRange broadcast as scalars
+# (fixes issue with SymbolicIndexingInterface where broadcasting over RaggedEnd would fail)
+ragged_idx = lastindex(ragged, 1)
+@test ragged_idx isa RecursiveArrayTools.RaggedEnd
+@test Base.broadcastable(ragged_idx) isa Ref
+# Broadcasting with RaggedEnd should not error
+@test identity.(ragged_idx) === ragged_idx
+
+ragged_range_idx = 1:lastindex(ragged, 1)
+@test ragged_range_idx isa RecursiveArrayTools.RaggedRange
+@test Base.broadcastable(ragged_range_idx) isa Ref
+# Broadcasting with RaggedRange should not error
+@test identity.(ragged_range_idx) === ragged_range_idx
+
 # Broadcasting of heterogeneous arrays (issue #454)
 u = VectorOfArray([[1.0], [2.0, 3.0]])
 @test length(view(u, :, 1)) == 1
@@ -220,8 +234,8 @@ u[1, [1, 3], 2] .= [7.0, 9.0]
 
 # 3D inner arrays (tensors) with ragged third dimension
 u = VectorOfArray([zeros(2, 1, n) for n in (2, 3)])
-@test size(view(u,:,:,:,1)) == (2, 1, 2)
-@test size(view(u,:,:,:,2)) == (2, 1, 3)
+@test size(view(u, :, :, :, 1)) == (2, 1, 2)
+@test size(view(u, :, :, :, 2)) == (2, 1, 3)
 # assign into a slice of the second inner array using last index Int
 u[2, 1, :, 2] .= [7.0, 8.0, 9.0]
 @test vec(u.u[2][2, 1, :]) == [7.0, 8.0, 9.0]