Fix fill!(::ArrayPartition)

src/array_partition.jl

@@ -209,13 +209,19 @@ function Base.copyto!(A::ArrayPartition, src::ArrayPartition)
     A
 end
 
+function Base.fill!(A::ArrayPartition, x)
+    unrolled_foreach!(A.x) do x_
+        fill!(x_, x)
+    end
+    A
+end
+
 function recursivefill!(b::ArrayPartition, a::T2) where {T2 <: Union{Number, Bool}}
     unrolled_foreach!(b.x) do x
         fill!(x, a)
     end
 end
 
-
 ## indexing
 
 # Interface for the linear indexing. This is just a view of the underlying nested structure

test/gpu/arraypartition_gpu.jl

@@ -14,3 +14,7 @@ mask = pA .> 0
 # Test recursive filling is done using GPU kernels and not scalar indexing
 RecursiveArrayTools.recursivefill!(pA, true)
 @test all(pA .== true)
+
+# Test that regular filling is done using GPU kernels and not scalar indexing
+fill!(pA, false)
+@test all(pA .== false)