Support global sparsity tracing through distribution constructors

Project.toml

@@ -1,7 +1,7 @@
 name = "Distributions"
 uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
 authors = ["JuliaStats"]
-version = "0.25.125"
+version = "0.25.126"
 
 [deps]
 AliasTables = "66dad0bd-aa9a-41b7-9441-69ab47430ed8"
@@ -23,11 +23,13 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [extensions]
 DistributionsChainRulesCoreExt = "ChainRulesCore"
 DistributionsDensityInterfaceExt = "DensityInterface"
+DistributionsSparseConnectivityTracerExt = "SparseConnectivityTracer"
 DistributionsTestExt = "Test"
 
 [compat]
@@ -50,6 +52,7 @@ Printf = "<0.0.1, 1"
 QuadGK = "2"
 Random = "<0.0.1, 1"
 SparseArrays = "<0.0.1, 1"
+SparseConnectivityTracer = "1"
 SpecialFunctions = "1.2, 2"
 StableRNGs = "1"
 StaticArrays = "1"
@@ -73,9 +76,10 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "StableRNGs", "Calculus", "ChainRulesCore", "ChainRulesTestUtils", "DensityInterface", "Distributed", "FiniteDifferences", "ForwardDiff", "JSON", "Optim", "SparseArrays", "StaticArrays", "Test", "OffsetArrays"]
+test = ["Aqua", "StableRNGs", "Calculus", "ChainRulesCore", "ChainRulesTestUtils", "DensityInterface", "Distributed", "FiniteDifferences", "ForwardDiff", "JSON", "Optim", "SparseArrays", "SparseConnectivityTracer", "StaticArrays", "Test", "OffsetArrays"]

ext/DistributionsSparseConnectivityTracerExt.jl

@@ -0,0 +1,16 @@
+module DistributionsSparseConnectivityTracerExt
+
+using Distributions: Distributions
+using SparseConnectivityTracer: AbstractTracer
+
+# Only the global `TracerSparsityDetector` is problematic for `@check_args`: it uses primal-free
+# tracers (`GradientTracer`/`HessianTracer`, the subtypes of `AbstractTracer`), so the validity
+# comparisons (e.g. `σ ≥ 0`) return a tracer that cannot be used in boolean context. Marking these
+# as not checkable skips validation so distribution constructors can still be traced.
+#
+# Local sparsity detection (`TracerLocalSparsityDetector`) passes `Dual` values, which are `<: Real`
+# (not `<: AbstractTracer`) and carry a primal, so the comparisons evaluate to `Bool` as usual — the
+# default `_arg_checkable` applies there and the checks run normally.
+Distributions._arg_checkable(::AbstractTracer) = false
+
+end

src/utils.jl

@@ -39,11 +39,14 @@ macro check_args(D, setup_or_check, checks...)
         checks = (setup_or_check, checks...)
     end
 
-    # Generate expressions for each condition
+    # Generate expressions for each condition, collecting the explicitly-named arguments so they can
+    # be forwarded to `check_args` (used to decide whether the checks can be evaluated at all).
+    args_exprs = Any[]
     conds_exprs = map(checks) do check
         if Meta.isexpr(check, :tuple, 3)
             # argument, condition, and message specified
             arg = check.args[1]
+            push!(args_exprs, esc(arg))
             cond = check.args[2]
             message = string(D, ": ", check.args[3])
             return :(($(esc(cond))) || throw(DomainError($(esc(arg)), $message)))
@@ -57,6 +60,7 @@ macro check_args(D, setup_or_check, checks...)
             else
                 # only argument and condition specified
                 arg = check.args[1]
+                push!(args_exprs, esc(arg))
                 cond = cond_or_message
                 message = string(D, ": the condition ", cond, " is not satisfied.")
                 return :(($(esc(cond))) || throw(DomainError($(esc(arg)), $message)))
@@ -73,11 +77,24 @@ macro check_args(D, setup_or_check, checks...)
         Distributions.check_args($(esc(:check_args))) do
             $(__source__)
             $(setup_stmts...)
-            $(conds_exprs...)
+            if all(Distributions._arg_checkable, ($(args_exprs...),))
+                $(conds_exprs...)
+            end
         end
     end
 end
 
+#=
+    _arg_checkable(x)::Bool
+
+Whether a constructor argument `x` supports `@check_args` validation. Defaults to `true`. Package
+extensions may specialize this to `false` for number-like types whose comparisons cannot be used in
+boolean context — for example sparsity-detection tracers, whose primal-free comparisons (`σ ≥ 0`)
+return a tracer rather than a `Bool`. When any checked argument is not checkable, `@check_args` skips
+the checks so distribution constructors can still be traced.
+=#
+_arg_checkable(@nospecialize(_)) = true
+
 """
     check_args(f, check::Bool)
 

test/runtests.jl

@@ -32,6 +32,7 @@ const tests = [
     "multivariate/mvlognormal",
     "types", # extra file compared to /src
     "utils",
+    "sparseconnectivitytracer", # extra file compared to /src
     "samplers",
     "univariate/discrete/categorical",
     "univariates",

test/sparseconnectivitytracer.jl

@@ -0,0 +1,43 @@
+using Distributions
+using SparseConnectivityTracer
+using Test
+
+@testset "SparseConnectivityTracer" begin
+    @testset "global gradient sparsity traces through construction" begin
+        detector = TracerSparsityDetector()
+        # `mean(Normal(μ, σ)) == μ` depends only on the first argument, `std` only on the second.
+        @test jacobian_sparsity(x -> mean(Normal(x[1], x[2])), [1.0, 2.0], detector) ==
+            [true false]
+        @test jacobian_sparsity(x -> std(Normal(x[1], x[2])), [1.0, 2.0], detector) ==
+            [false true]
+    end
+
+    @testset "global hessian sparsity traces through construction" begin
+        detector = TracerSparsityDetector()
+        # Two independent `Normal`s. Within each, `mean(d) * std(d) == μσ` couples that
+        # distribution's own (μ, σ); the two distributions do not interact. So the Hessian is
+        # block off-diagonal.
+        f =
+            x ->
+                mean(Normal(x[1], x[2])) * std(Normal(x[1], x[2])) +
+                mean(Normal(x[3], x[4])) * std(Normal(x[3], x[4]))
+        @test hessian_sparsity(f, [1.0, 2.0, 3.0, 4.0], detector) == [
+            false true false false
+            true false false false
+            false false false true
+            false false true false
+        ]
+    end
+
+    @testset "local sparsity detection still validates arguments" begin
+        detector = TracerLocalSparsityDetector()
+        @test jacobian_sparsity(x -> mean(Normal(x[1], x[2])), [1.0, 2.0], detector) ==
+            [true false]
+        # The primal is available, so the `σ ≥ 0` check is enforced.
+        @test_throws Exception jacobian_sparsity(
+            x -> std(Normal(x[1], x[2])),
+            [1.0, -1.0],
+            detector,
+        )
+    end
+end