fix(overload): MissingPrimalError in reductions of tracer Fills with FillArrays v1.17

CHANGELOG.md

@@ -1,6 +1,7 @@
 # SparseConnectivityTracer.jl
 
 ## Version `v1.2.3`
+* ![Bugfix][badge-bugfix] Fix `MissingPrimalError` in reductions of tracer `Fill`s with FillArrays.jl 1.17 ([#328])
 * ![Maintenance][badge-maintenance] Allow LogExpFunctions.jl 1.x ([#327])
 
 ## Version `v1.2.2`
@@ -172,6 +173,7 @@ This release is only breaking for users who touched unexported internals.
 [badge-maintenance]: https://img.shields.io/badge/maintenance-gray.svg
 [badge-docs]: https://img.shields.io/badge/docs-orange.svg
 
+[#328]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/328
 [#327]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/327
 [#322]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/322
 [#292]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/292

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseConnectivityTracer"
 uuid = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 authors = ["Adrian Hill <gh@adrianhill.de>"]
-version = "1.2.2"
+version = "1.2.3-DEV"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

benchmark/Project.toml

@@ -2,6 +2,7 @@
 ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

benchmark/jacobian.jl

@@ -4,7 +4,8 @@ using SparseConnectivityTracerBenchmarks.ODE: Brusselator!, brusselator_2d_loop!
 
 using Random: MersenneTwister
 using SparseArrays: sprand
-using SimpleDiffEq: ODEProblem, solve, SimpleEuler
+using CommonSolve: solve
+using SimpleDiffEq: ODEProblem, SimpleEuler
 using Flux: Conv
 
 function jacbench(detector)

src/SparseConnectivityTracer.jl

@@ -7,7 +7,7 @@ using Random: AbstractRNG, SamplerType
 
 using LinearAlgebra: LinearAlgebra, Symmetric
 using LinearAlgebra: Diagonal, diag, diagind
-using FillArrays: Fill
+using FillArrays: Fill, AbstractFill, getindex_value
 
 using DocStringExtensions: DocStringExtensions, TYPEDEF, TYPEDFIELDS
 
@@ -25,6 +25,7 @@ include("overloads/three_arg.jl")
 include("overloads/ifelse_global.jl")
 include("overloads/is_functions.jl")
 include("overloads/arrays.jl")
+include("overloads/fillarrays.jl")
 include("overloads/ambiguities.jl")
 
 include("trace_functions.jl")

src/overloads/fillarrays.jl

@@ -0,0 +1,28 @@
+#===============#
+# FillArrays.jl #
+#===============#
+
+# Since v1.17, FillArrays' reduction fast paths call `iszero`/`isone` on the fill value,
+# which global tracers deliberately answer with a `MissingPrimalError`,
+# so reductions of tracer `Fill`s get exact O(1) overloads.
+
+function Base.sum(A::AbstractFill{T}; dims = :) where {T <: AbstractTracer}
+    t = isempty(A) ? myempty(T) : getindex_value(A)
+    dims isa Colon && return t
+    return Fill(t, Base.reduced_indices(axes(A), dims))
+end
+
+# `second_order_or` degrades to a plain union on `GradientTracer`s,
+# so a single method is exact for both tracer types.
+function Base.prod(A::AbstractFill{T}; dims = :) where {T <: AbstractTracer}
+    t = getindex_value(A)
+    if dims isa Colon
+        isempty(A) && return myempty(T)
+        length(A) == 1 && return t
+        return second_order_or(t, t)
+    end
+    ri = Base.reduced_indices(axes(A), dims)
+    isempty(A) && return Fill(myempty(T), ri)
+    length(A) == prod(length, ri) && return Fill(t, ri) # one element per output slice
+    return Fill(second_order_or(t, t), ri)
+end

test/Project.toml

@@ -4,10 +4,12 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"

test/test_arrays.jl

@@ -9,6 +9,7 @@ using SparseConnectivityTracer:
     tracer
 using Test
 
+using FillArrays: Fill
 using LinearAlgebra: Symmetric, Diagonal, diagind
 using LinearAlgebra: det, logdet, logabsdet, norm, opnorm
 using LinearAlgebra: eigen, eigmax, eigmin
@@ -362,6 +363,91 @@ end
     end
 end
 
+@testset "Fill reductions" begin
+    J(f, x) = jacobian_sparsity(f, x, detector)
+    H(f, x) = hessian_sparsity(f, x, detector)
+    x = rand(2)
+
+    @testset "sum" begin
+        @test J(x -> sum(Fill(x[1] + x[2], 3)), x) == [1 1]
+        # Summation is linear: the pattern of the fill value is preserved exactly
+        @test H(x -> sum(Fill(x[1] + x[2], 3)), x) == [0 0; 0 0]
+        @test H(x -> sum(Fill(x[1] * x[2], 3)), x) == [0 1; 1 0]
+        # Reductions along `dims` preserve the pattern
+        # (that they return another `Fill` is tested in "Fill representation is preserved")
+        @test J(x -> sum(sum(Fill(x[1] + x[2], 3, 4); dims = 2)), x) == [1 1]
+        # Empty sums have no dependencies
+        @test J(x -> sum(Fill(x[1] + x[2], 0)), x) == [0 0]
+        @test H(x -> sum(Fill(x[1] * x[2], 0)), x) == [0 0; 0 0]
+    end
+
+    @testset "prod" begin
+        @test J(x -> prod(Fill(x[1] + x[2], 3)), x) == [1 1]
+        # Products create second-order self-interactions...
+        @test H(x -> prod(Fill(x[1] + x[2], 3)), x) == [1 1; 1 1]
+        # ...except on a single element, which is returned exactly
+        @test H(x -> prod(Fill(x[1] + x[2], 1)), x) == [0 0; 0 0]
+        @test H(x -> prod(Fill(x[1] * x[2], 1)), x) == [0 1; 1 0]
+        # Reductions along `dims` reduce over the lengths of these dimensions
+        @test H(x -> sum(prod(Fill(x[1] + x[2], 3, 4); dims = 1)), x) == [1 1; 1 1]
+        @test H(x -> sum(prod(Fill(x[1] + x[2], 3, 4); dims = 3)), x) == [0 0; 0 0]
+        # Empty products have no dependencies
+        @test J(x -> prod(Fill(x[1] + x[2], 0)), x) == [0 0]
+    end
+
+    @testset "Fill representation is preserved" begin
+        t = idx2tracer([1, 2])
+        A = Fill(t, 3, 4)
+
+        @test sum(A) isa TG
+        @test sameidx(sum(A), [1, 2])
+        @test sameidx(prod(A), [1, 2])
+
+        # Explicit `dims = :` behaves like the default: a scalar tracer, not a `Fill`,
+        # with the same output dimensionality as reductions of regular arrays
+        @test sum(A; dims = :) isa TG
+        @test prod(A; dims = :) isa TG
+        @test ndims(sum(A; dims = :)) == ndims(sum(ones(3, 4); dims = :))
+        @test ndims(prod(A; dims = :)) == ndims(prod(ones(3, 4); dims = :))
+
+        S = sum(A; dims = 2)
+        @test S isa Fill
+        @test size(S) == (3, 1)
+        @test sameidx(S[1, 1], [1, 2])
+
+        P = prod(A; dims = 1)
+        @test P isa Fill
+        @test size(P) == (1, 4)
+        @test sameidx(P[1, 1], [1, 2])
+
+        # Tuple and vector `dims`
+        for dims in ((1, 2), [1, 2])
+            S12 = sum(A; dims = dims)
+            @test S12 isa Fill
+            @test size(S12) == (1, 1)
+            @test sameidx(S12[1, 1], [1, 2])
+
+            P12 = prod(A; dims = dims)
+            @test P12 isa Fill
+            @test size(P12) == (1, 1)
+        end
+
+        # `dims` beyond `ndims(A)` reduce over a single element, so the tracer is exact
+        S3 = sum(A; dims = 3)
+        @test S3 isa Fill
+        @test size(S3) == (3, 4)
+        @test sameidx(S3[1, 1], [1, 2])
+
+        P3 = prod(A; dims = 3)
+        @test P3 isa Fill
+        @test size(P3) == (3, 4)
+        @test sameidx(P3[1, 1], [1, 2])
+
+        @test isemptytracer(sum(Fill(t, 0)))
+        @test isemptytracer(prod(Fill(t, 0)))
+    end
+end
+
 @testset "clamp!" begin
     t1 = idx2tracer(1)
     t2 = idx2tracer(2)