Inline optimized any/all for ArrayPartition on Julia 1.13+

Julia 1.13 (JuliaLang/julia#61184) removes the f::Function restriction
from any/all, so defining any(f, ::ArrayPartition) no longer causes ~780
invalidations. On 1.13+, the optimized partition-level methods are defined
directly in the main package. On older Julia, the GPU-check fallback
remains and the RecursiveArrayToolsArrayPartitionAnyAll subpackage is
needed for optimized methods.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

src/array_partition.jl

@@ -741,37 +741,56 @@ ODEProblem(func, AP[ [1.,2.,3.], [1. 2.;3. 4.] ], (0, 1)) |> solve
 """
 struct AP end
 
-# any/all: provide methods that work partition-by-partition to avoid scalar indexing on GPUs.
-# Without RecursiveArrayToolsArrayPartitionAnyAll, the AbstractArray fallback iterates
-# element-by-element, which triggers scalar GPU indexing errors. These methods check for
-# GPU sub-arrays and give a helpful error, or fall through to the partition-level implementation.
-const _ANYALL_HINT = """
-    `any`/`all` on `ArrayPartition` with GPU arrays requires loading the subpackage:
-        using RecursiveArrayToolsArrayPartitionAnyAll
-    This provides optimized partition-level `any`/`all` that avoids scalar GPU indexing.
-    """
-
-function _check_gpu_anyall(A::ArrayPartition)
-    for x in A.x
-        if x isa GPUArraysCore.AnyGPUArray
-            error(_ANYALL_HINT)
+# any/all on ArrayPartition: partition-level iteration avoids scalar GPU indexing
+# and is ~1.5-1.8x faster than element-by-element AbstractArray fallback.
+#
+# On Julia ≥ 1.13, Base no longer restricts any/all to f::Function, so defining
+# any(f, ::ArrayPartition) causes minimal invalidation. On older Julia, the
+# f::Function methods caused ~780 invalidations, so the optimized methods were
+# separated into RecursiveArrayToolsArrayPartitionAnyAll.
+#
+# We now define them inline with a GPU check: if sub-arrays are GPU arrays and
+# we're on old Julia without the extension, give a helpful error.
+@static if VERSION >= v"1.13.0-DEV.0"
+    # Julia 1.13+: safe to define directly, minimal invalidation
+    Base.any(f, A::ArrayPartition) = any((any(f, x) for x in A.x))
+    Base.any(f::Function, A::ArrayPartition) = any((any(f, x) for x in A.x))
+    Base.any(A::ArrayPartition) = any(identity, A)
+    Base.all(f, A::ArrayPartition) = all((all(f, x) for x in A.x))
+    Base.all(f::Function, A::ArrayPartition) = all((all(f, x) for x in A.x))
+    Base.all(A::ArrayPartition) = all(identity, A)
+else
+    # Julia < 1.13: only define GPU-check methods that error for GPU arrays
+    # and fall through to AbstractArray for CPU. The optimized partition-level
+    # methods are in RecursiveArrayToolsArrayPartitionAnyAll to avoid invalidations.
+    const _ANYALL_GPU_HINT = """
+        `any`/`all` on `ArrayPartition` with GPU arrays requires loading the subpackage:
+            using RecursiveArrayToolsArrayPartitionAnyAll
+        This provides optimized partition-level `any`/`all` that avoids scalar GPU indexing.
+        """
+
+    function _check_gpu_anyall(A::ArrayPartition)
+        for x in A.x
+            if x isa GPUArraysCore.AnyGPUArray
+                error(_ANYALL_GPU_HINT)
+            end
         end
     end
-end
 
-function Base.any(f::Function, A::ArrayPartition)
-    _check_gpu_anyall(A)
-    return Base.invoke(any, Tuple{Function, AbstractArray}, f, A)
-end
-function Base.all(f::Function, A::ArrayPartition)
-    _check_gpu_anyall(A)
-    return Base.invoke(all, Tuple{Function, AbstractArray}, f, A)
-end
-function Base.any(f, A::ArrayPartition)
-    _check_gpu_anyall(A)
-    return Base.invoke(any, Tuple{Any, AbstractArray}, f, A)
-end
-function Base.all(f, A::ArrayPartition)
-    _check_gpu_anyall(A)
-    return Base.invoke(all, Tuple{Any, AbstractArray}, f, A)
+    function Base.any(f::Function, A::ArrayPartition)
+        _check_gpu_anyall(A)
+        return Base.invoke(any, Tuple{Function, AbstractArray}, f, A)
+    end
+    function Base.all(f::Function, A::ArrayPartition)
+        _check_gpu_anyall(A)
+        return Base.invoke(all, Tuple{Function, AbstractArray}, f, A)
+    end
+    function Base.any(f, A::ArrayPartition)
+        _check_gpu_anyall(A)
+        return Base.invoke(any, Tuple{Any, AbstractArray}, f, A)
+    end
+    function Base.all(f, A::ArrayPartition)
+        _check_gpu_anyall(A)
+        return Base.invoke(all, Tuple{Any, AbstractArray}, f, A)
+    end
 end

test/partitions_test.jl

@@ -178,16 +178,21 @@ recursivecopy!(dest_ap, src_ap)
 @inferred mapreduce(string, *, x)
 @test mapreduce(i -> string(i) * "q", *, x) == "1q2q3.0q4.0q"
 
-# any/all — the base package defines GPU-aware methods that error for GPU sub-arrays
-# and fall through to AbstractArray iteration for CPU arrays. The optimized partition-level
-# iteration (1.5-1.8x faster) requires `using RecursiveArrayToolsArrayPartitionAnyAll`.
-# These @test_broken verify the optimized extension is NOT loaded.
-@test_broken occursin("ArrayPartitionAnyAll",
-    string(which(any, Tuple{Function, ArrayPartition}).module))
-@test_broken occursin("ArrayPartitionAnyAll",
-    string(which(all, Tuple{Function, ArrayPartition}).module))
-
-# Correctness tests pass via the base package methods (element-by-element on CPU)
+# any/all — on Julia ≥ 1.13, optimized partition-level methods are defined inline.
+# On older Julia, they require `using RecursiveArrayToolsArrayPartitionAnyAll`.
+@static if VERSION >= v"1.13.0-DEV.0"
+    # Optimized methods are always active on 1.13+
+    @test which(any, Tuple{Function, ArrayPartition}).module === RecursiveArrayTools
+    @test which(all, Tuple{Function, ArrayPartition}).module === RecursiveArrayTools
+else
+    # On older Julia, optimized methods require the extension
+    @test_broken occursin("ArrayPartitionAnyAll",
+        string(which(any, Tuple{Function, ArrayPartition}).module))
+    @test_broken occursin("ArrayPartitionAnyAll",
+        string(which(all, Tuple{Function, ArrayPartition}).module))
+end
+
+# Correctness tests
 @test !any(isnan, AP[[1, 2], [3.0, 4.0]])
 @test !any(isnan, AP[[3.0, 4.0]])
 @test  any(isnan, AP[[NaN], [3.0, 4.0]])