diff --git a/CHANGELOG.md b/CHANGELOG.md index 68a6f0a3..560ff057 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/Project.toml b/Project.toml index 870b0754..5d78602e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseConnectivityTracer" uuid = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" authors = ["Adrian Hill "] -version = "1.2.2" +version = "1.2.3-DEV" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/benchmark/Project.toml b/benchmark/Project.toml index e1388044..7767788f 100644 --- a/benchmark/Project.toml +++ b/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" diff --git a/benchmark/jacobian.jl b/benchmark/jacobian.jl index 09a70f35..4b349ba6 100644 --- a/benchmark/jacobian.jl +++ b/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) diff --git a/src/SparseConnectivityTracer.jl b/src/SparseConnectivityTracer.jl index 25bce8eb..2335c632 100644 --- a/src/SparseConnectivityTracer.jl +++ b/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") diff --git a/src/overloads/fillarrays.jl b/src/overloads/fillarrays.jl new file mode 100644 index 00000000..431d8c88 --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index c106f3f6..6788e9a3 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/test_arrays.jl b/test/test_arrays.jl index bbba8acf..289f6a3e 100644 --- a/test/test_arrays.jl +++ b/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)