[sethaxen]

This package removes a number of dependencies that we primarily only needed for a few utilities that are easily reproduced here. In return, on my machine, package load time is reduced from 0.36s to 0.05s 0.04s 0.25s in an interactive session and from .5s to .27s in a non-interactive session.

Making this package lighter weight allows it to be a dependency of other lightweight packages, e.g. PSIS, which currently has a load time of 0.01s.

Distributions

This is our heaviest dependency, and in most places we just use it to call StatsFuns functions anyways, so now we call StatsFuns. In a few places we sample from a categorical distribution. That can be implemented efficiently in 14 LOC, so now we roll that ourselves. The only place where Distributions integration is required is when passing a probabilistic sampler to rstar, which then returns a PoissonBinomial. In this PR, that has been moved to an extension, and the user is now required to load Distributions if they want to use a probabilistic classifier. That being said, since MLJBase depends on Distributions, in most cases users won't need to do that. Nevertheless, this makes this a breaking change.

StatsBase

We just needed wsample (implemented above), tiedrank, and counts, which again require just 22 and 10 LOC, respectively, so they are included here.

DataStructures

SortedDict was never really necessary but just allowed just to remove a sort call at the end of unique_indices. Previously DataStructures was an indirect dependency, but with the above 2 removals, it no longer is, so this PR also removes it. and it still is via StatsBase, but this PR removes it as a direct dependency.

MLJModelInterface/Tables

While MMI is itself lightweight, in noninteractive sessions it pulled in REPL and increased package load time by about 25%. While how we use MMI is the intended behavior, it is useless apart from MLJBase which contains its implementations. This split predates the introduction of package extensions, and it seems more natural here to have an MLJBase extensions so that we can pin compat to MLJBase versions that include Tables and Distributions as dependencies (a breaking change to MLJBase would be required for it to drop these deps). That allows rstar to return a Distributions object when passed a probabilistic classifier with absolutely no deviation from current behavior and therefore no breaking release required.

[github-actions]

MCMCDiagnosticTools.jl documentation for PR #177 is available at:
https://TuringLang.github.io/MCMCDiagnosticTools.jl/previews/PR177/

[codecov]

Codecov Report

❌ Patch coverage is 98.85057% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 97.05%. Comparing base (559b68e) to head (130dc81).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/discretediag.jl	95.23%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #177      +/-   ##
==========================================
+ Coverage   97.01%   97.05%   +0.04%     
==========================================
  Files          11       13       +2     
  Lines         937      952      +15     
==========================================
+ Hits          909      924      +15     
  Misses         28       28              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[coveralls]

Coverage Report for CI Build 40316420566

Warning

No base build found for commit 69590c2 on main.
Coverage changes can't be calculated without a base build.
If a base build is processing, this comment will update automatically when it completes.

Coverage: 97.059%

Details

• Patch coverage: 1 uncovered change across 1 file (86 of 87 lines covered, 98.85%).

Uncovered Changes

File	Changed	Covered	%
src/discretediag.jl	21	20	95.24%
Total (7 files)	87	86	98.85%

Coverage Regressions

Requires a base build to compare against. How to fix this →

---

Coverage Stats

Relevant Lines:	952
Covered Lines:	924
Line Coverage:	97.06%
Coverage Strength:	1139018.93 hits per line

---

💛 - Coveralls

[sethaxen]

In practice this is unreachable, because MMI just defines the interface functions we call but they are implemented in MLJBase (see https://github.com/JuliaAI/MLJModelInterface.jl/blob/f18e0e5692327eda5352caa1430bd5e938898dbe/src/data_utils.jl#L2-L15), which depends on Distributions. So if the user did not yet load MLJBase, they would have already errored before getting to _rstar_distribution, and if they did already load MLJBase, then Distributions would be in the environment, as would our package extension, and _rstar_distribution would not throw a method error.

However, MLJBase could make a breaking release dropping Distributions as a direct dependency without MLJModelInterface requiring a breaking release, so this error hint should still exist for that hypothetical future.

[sethaxen]

Another breaking option would be to define our own PoissonBinomial object and return that, documenting only its fields and that it can be converted to a Distributions object, which could be in the extension. Then no error is ever needed.

[sethaxen]

Yet another option, since MMI can't do anything until MLJBase is loaded anyways, would be to move the rstar implementations to an extension weakly depending on either MMI or MLJBase directly. The latter allows us to forego any requirement that users load Distributions, since then we can specify via the compat MLJBase versions that have Distributions as a direct dependency. While MMI is lightweight, removing it and Tables as a direct dependency reduces package load time on my machine from .2s in a script to .06s, I assume because MMI depends on and loads REPL.

[devmotion]

I think it's generally good to drop unnnecessary dependencies (as DataStructures seems to be) and move optional functionality to extensions if it doesn't break typical workflows or makes them significantly less convenient.

I also think that it's a bit simpler to directly use StatsFuns functions instead of going through Distributions.

However, usually I think it's a bad sign if one starts to reimplement functionality available somewhere else, which seems to be the case with the StatsBase functionality we are using. I would also assume that most people using MCMCDiagnosticTools already have StatsBase (and likely also Distributions?) available in their environment anyway, so load time of these packages cannot be attributed solely to MCMCDiagnosticTools. Additionally, <0.4 seconds seems already rather fast.

How much of the load time can be attributed to each dependency (in an environment with only MCMCDiagnosticTools), did you check it with @time_using? Moreover, I wonder if we could make use of https://julialang.github.io/PrecompileTools.jl to reduce TTFX?

As a side remark, I recently learnt that __init__ can slow down load times quite a bit: JuliaMath/InverseFunctions.jl#67

(A second side remark: Skimming through the changes made me wonder whether simulate_DAR1! would benefit from the alias-sampler functionality in StatsBase - which is based on AliasTables)

[sethaxen]

(A second side remark: Skimming through the changes made me wonder whether simulate_DAR1! would benefit from the alias-sampler functionality in StatsBase - which is based on AliasTables)

Agreed, I had thought so as well. AliasTables adds almost nothing to the load time, so I've switched to it here. Also, for typical values of nsim, the last remaining call to wsample on main made discretediag much slower than using an alias table, so I switched to using it there as well.

How much of the load time can be attributed to each dependency (in an environment with only MCMCDiagnosticTools), did you check it with @time_using?

Yes I did. @time_imports using MCMCDiagnosticTools on main (Julia 1.13) shows the main contributors to load time as

Package	Pulled in by	Non-interactive	Interactive
Total using time		569 ms	439 ms
SparseArrays	StatsBase, PDMats (Distributions)	167 ms	181 ms
REPL	MLJModelInterface	95 ms	already loaded
FillArrays	Distributions	39 ms	70 ms
StyledStrings	REPL/Markdown (MLJModelInterface)	25 ms	already loaded
Dates	TOML (via SpecialFunctions), Accessors (via Distributions)	25 ms	already loaded
Distributions	direct	24 ms	25 ms
AbstractFFTs	direct (~90% compilation)	21 ms	19 ms
DataStructures	direct, StatsBase	12 ms	12 ms
Markdown	REPL/InteractiveUtils (MLJModelInterface)	11 ms	already loaded
OpenSpecFun_jll	SpecialFunctions (~96% compilation)	10 ms	10 ms
StatsBase	direct	7 ms	8 ms
Tables	direct	6 ms	6 ms
JuliaSyntaxHighlighting	REPL/Markdown (MLJModelInterface)	3 ms	already loaded
MLJModelInterface	direct	2 ms	2 ms

So the biggest contributors in interactive sessions are SparseArrays, FillArrays, and Distributions itself. StatsBase itself doesn't contribute much to load time, but by pulling in DataStructures and SparseArrays, it accounts for about .2s, or almost half of the package load time. In non-interactive sessions, MMI indirectly accounts for .14s, or about 25% of package load time.

However, usually I think it's a bad sign if one starts to reimplement functionality available somewhere else, which seems to be the case with the StatsBase functionality we are using.

I would typically agree with this when the functionality being reimplemented is sufficiently complex. In this case (with the recent commits), it's just _counts and _tiedrank!, which are very simple. To the point that if I started out with a brand new package with a .06s load time, I wouldn't add a StatsBase dependency just to use its versions of these functions if it increased the package load time by .2s. Because we know when these functions are being called, they don't need to be as general as those in StatsBase either, and this simplifies things a bit. e.g. _tiedrank! can reuse an already-allocated container. This causes small speed-ups but a 25% reduction in allocations in an inner loop, e.g.

using Random, MCMCDiagnosticTools, BenchmarkTools
rng = Xoshiro(42)
draws = randn(rng, 1_000, 4, 1_000)
@btime ess_rhat($draws);

main:

  461.233 ms (87632 allocations: 313.72 MiB)

here:

  449.247 ms (65635 allocations: 251.82 MiB)

I would also assume that most people using MCMCDiagnosticTools already have StatsBase (and likely also Distributions?) available in their environment anyway, so load time of these packages cannot be attributed solely to MCMCDiagnosticTools. Additionally, <0.4 seconds seems already rather fast.

It's hard to know about users, but 18 registered packages depend directly on MCMCDT. All of them also depend on StatsBase, and 16 depend on Distributions (ParetoSmooth and FlexiChains after this PR would no longer indirectly import Distributions, which I think is good). While I agree that most users will already have StatsBase in their envs, it's certainly not all. Historically, Soss/Tilde tried to get rid of Distributions as a dependency. StanBlocks.jl builds probabilistic models ready-for-sampling without depending on either. Other PPLs that opt out of Distributions/StatsBase could want to include our diagnostics as well.

But there's actually already an example of a package that wants MCMCDT w/o StatsBase. PSIS is intentionally extremely lightweight to encourage any package that does importance sampling or Monte Carlo estimation to use its diagnostics internally, regardless of their stack. In the convergence diagnostics coming shortly, it really should by default use relative ESS to avoid silently doing the wrong thing for MCMC draws, but currently it wouldn't because it would increase the load time by ~30x. Agreed that <0.4s is fast, but more than once, packages with quicker load times have replicated functionality from small packages I've maintained to avoid increasing by a few hundred ms.

As a side remark, I recently learnt that __init__ can slow down load times quite a bit: JuliaMath/InverseFunctions.jl#67

Interesting, I had not seen that. I attempted to replicate that here but did not see any slowdown whether we used __init__ or not, so I've left it in for now.

I wonder if we could make use of https://julialang.github.io/PrecompileTools.jl to reduce TTFX?

Absolutely, and at least ess_rhat would be a good target for precompilation, but I think that should be orthogonal to this PR.

[sethaxen]

Following my suggestion in #177 (comment), I've also removed MLJModelInterface and Tables as dependencies and instead moved the rstar implementations to an extension triggered on MLJBase. See OP for a justification. As a result, this PR is now non-breaking.

[sethaxen]

@devmotion does this PR look fine now to you?

[devmotion]

IMO it's still strange to duplicate StatsBase functionality in this package, in particular since in contrast to MLJ/MLJBase I think in most applications it will be loaded anyway. How much of the load time can actually be attributed to StatsBase?

[sethaxen]

IMO it's still strange to duplicate StatsBase functionality in this package, in particular since in contrast to MLJ/MLJBase I think in most applications it will be loaded anyway. How much of the load time can actually be attributed to StatsBase?

After more time working on the PSIS use-case, I decided it actually doesn't make sense for it to depend on MCMCDiagnosticTools after all, so that need is removed. As a result, I've reverted the removal of StatsBase a dependency.

The load time improvements are now much more modest; down to .25s in an interactive session from the original .36s. Most of the remaining load time is from StatsBase. When it's removed as a dependency, load time is .04s.

[devmotion]

The StatsBase loadtime seems a bit extreme, do you know why it takes so much time to load? I wonder if we could improve it.

[sethaxen]

The StatsBase loadtime seems a bit extreme, do you know why it takes so much time to load? I wonder if we could improve it.

StatsBase's load time is mostly (~80%) due to it loading SparseArrays:

julia> @time using StatsBase
  0.221406 seconds (177.65 k allocations: 11.268 MiB)

julia> @time_imports using StatsBase
      0.5 ms  DataAPI
      4.9 ms  OrderedCollections
     15.5 ms  DataStructures
      0.8 ms  SortingAlgorithms
      3.2 ms  Missings
               ┌ 0.0 ms DocStringExtensions.__init__()
      2.9 ms  DocStringExtensions
      7.9 ms  IrrationalConstants
      0.6 ms  LogExpFunctions
      1.0 ms  Statistics
               ┌ 0.2 ms CompilerSupportLibraries_jll.__init__()
      0.7 ms  CompilerSupportLibraries_jll
               ┌ 0.7 ms SuiteSparse_jll.__init__()
      1.0 ms  SuiteSparse_jll
      1.0 ms  Serialization
    171.3 ms  SparseArrays
      0.4 ms  Statistics → SparseArraysExt
      0.4 ms  StatsAPI
      0.4 ms  PtrArrays
      1.2 ms  AliasTables
      8.9 ms  StatsBase

This issue has been raised several times in StatsBase (JuliaStats/StatsBase.jl#885 and JuliaStats/StatsBase.jl#943). I noticed you contributed to both discussions.

Claude Opus 5.5 noted that SparseArrays' long load time is not due to invalidation but rather due to spending a lot of time in jl_activate_methods, which updates existing method tables with a package's methods; SparseArrays adds about 750 methods to functions defined in Base/LinearAlgebra.

I also asked Claude to bring back the change JuliaStats/StatsBase.jl#885 and redo our loading times here with that version of StatsBase as a dependency, and it measured a .07s load time. So there's really nothing we can do or need to do here to get the load time any lower; it needs to happen upstream in StatsBase.

[devmotion]

Thanks, that matches what I found. I opened JuliaStats/StatsBase.jl#1016 (deprecates indicatormat(...; sparse=true) in favour of indicatormat(SparseMatrixCSC, x, ...), non-breaking) and JuliaStats/StatsBase.jl#1017 (moves SparseArrays to an extension for StatsBase 0.35). With #1017, using StatsBase drops from ~265 ms to ~95 ms on my machine.