From 709893b5c6d2672b0da2bc4dc049da70de90c953 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 17 Jun 2026 00:45:50 +0200 Subject: [PATCH] Fix flaky ChainRules AD tests by keeping finite differences off domain boundaries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AD tests in `test/chainrules.jl` intermittently failed in CI with `DomainError`s from `loggamma`/`digamma`/`log`. The analytic rules in the ChainRulesCore extension are correct; the failures came from the numerical reference: `test_frule`/`test_rrule` perturb the inputs with finite differences, and with unseeded `exp(randn())`/`logistic(randn())` draws an argument occasionally landed close to a domain boundary, so the stencil crossed it. The fix keeps every stencil in-domain by construction, relying on ChainRulesTestUtils' defaults rather than a custom `fdm` or tangents. `test_frule`/`test_rrule` evaluate the reference at `input ± reach` with `reach = max_range * |tangent|`, and with CRTU's defaults both factors are bounded: * `max_range = 1e-2` (CRTU's `_fdm = central_fdm(5, 1; max_range = 1e-2)`); * `Float64` tangents are drawn from `-9:0.01:9`, so `|tangent| <= 9`. Hence `reach <= 0.09`. Drawing every differentiated argument >= 0.1 from its boundary (`0.1 + randexp()`, `0.1 + 0.8rand()`) then makes a boundary crossing impossible (`0.09 < 0.1`), for any draw or seed; each log-pdf domain is an independent box constraint, so per-coordinate anchoring suffices. The wide margin also keeps the points well-conditioned, so the reference stays accurate at the default tolerance. A fixed seed is added for reproducibility, and integer arguments resolve to `NoTangent()` automatically. The margin relies on CRTU's defaults; a comment records that the `0.1` anchor must grow if `max_range` or the tangent range grows. No new dependency is added. Verified: 0 failures over 2500 seeds x 7 distributions (frule + rrule); full Pkg.test() passes (chainrules | 98 98, QA | 19 19). Co-Authored-By: Claude Opus 4.8 (1M context) --- test/chainrules.jl | 49 +++++++++++++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 18 deletions(-) diff --git a/test/chainrules.jl b/test/chainrules.jl index d59b5eac..03113e4b 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -4,42 +4,55 @@ using ChainRulesTestUtils using Random @testset "chainrules" begin - x = exp(randn()) - y = exp(randn()) - z = logistic(randn()) + # `test_frule`/`test_rrule` finite-difference the log-pdfs, which throw `DomainError` + # outside their domain. ChainRulesTestUtils evaluates the reference at `input ± reach` + # with `reach = max_range * |tangent|`. With its defaults (`max_range = 1e-2` and, for + # `Float64`, `|tangent| ≤ 9` since tangents are drawn from `-9:0.01:9`) the reach is at + # most `0.09`. Drawing every differentiated argument ≥ 0.1 from its boundary therefore + # keeps the stencil inside the domain by construction (`0.09 < 0.1`), for any draw or + # seed, so we can rely on the CRTU defaults without a custom `fdm` or tangents. This + # margin assumes those CRTU defaults; if `max_range` or the tangent range grows, the + # `0.1` anchor below must grow with it. A fixed seed makes any failure reproducible. + Random.seed!(1234) + pos() = 0.1 + randexp() # positive args (shape/df/rate/λ/eval point) + unit01() = 0.1 + 0.8rand() # args in (0, 1): beta `x`, binom `p` + + x = pos() + y = pos() + z = unit01() test_frule(betalogpdf, x, y, z) test_rrule(betalogpdf, x, y, z) - x = exp(randn()) - y = exp(randn()) - z = exp(randn()) + x = pos() + y = pos() + z = pos() test_frule(gammalogpdf, x, y, z) test_rrule(gammalogpdf, x, y, z) - x = exp(randn()) - y = exp(randn()) + x = pos() + y = pos() test_frule(chisqlogpdf, x, y) test_rrule(chisqlogpdf, x, y) - x = exp(randn()) - y = exp(randn()) - z = exp(randn()) + x = pos() + y = pos() + z = pos() test_frule(fdistlogpdf, x, y, z) test_rrule(fdistlogpdf, x, y, z) - x = exp(randn()) - y = randn() + x = pos() + y = randn() # location: unbounded test_frule(tdistlogpdf, x, y) test_rrule(tdistlogpdf, x, y) - x = rand(1:100) - y = logistic(randn()) - z = rand(1:x) + x = rand(1:100) # `n`: integer, NoTangent + y = unit01() + z = rand(1:x) # `k`: integer, NoTangent test_frule(binomlogpdf, x, y, z) test_rrule(binomlogpdf, x, y, z) - x = exp(randn()) - y = rand(1:100) + x = pos() + y = rand(1:100) # `k`: integer, NoTangent test_frule(poislogpdf, x, y) test_rrule(poislogpdf, x, y)