Reduce number of dependencies

[sethaxen]

@ParadaCarleton and I discussed trying to slim down the package to make it a more lightweight dependency. Currently on my machine this package loads in 5s, vs 5ms with https://github.com/arviz-devs/PSIS.jl, which only depends on standard library packages.

The two reasons to add dependencies are 1) convenience and 2) performance. In my view, features like pretty-printed results are not worth the extra load time if a package can be light-weight, or they can be made conditional dependencies.

For performance, we discussed that Tullio+LoopVectorization may be unnecessary, as PSIS.jl is much faster than ParetoSmooth.jl for vectors and only 20% slower than ParetoSmooth.jl for arrays (if PSIS.jl uses the @threads macro like ParetoSmooth.jl does). If removing then only causes a ~20% decrease in performance but decreases load time significantly, this seems acceptable.

[sethaxen]

On my machine, a non-exhaustive list of load times of the heaviest dependencies in a clean environment:

• LoopVectorization: 2.7s
• AxisKeys: 0.86s
• MCMCDiagnosticTools: 0.80s
• Distributions: 0.77s
• Polyester: 0.60 s

[ParadaCarleton]

Yep, makes sense. I think a fair bit of these can be cut out.

MCMCDiagnosticTools, as we discussed earlier, can be cut out pretty easily. Another way we could cut it out besides what we discussed is just copy/pasting the FFTESS code from MCMCDiagnosticTools so we don't need to take everything else, which might just be the easiest option. I think both Distributions.jl and Polyester.jl can be cut out too at very low cost and I'll make a PR doing that soon.

It might be worth checking whether switching from AxisKeys to DimensionalData reduces compile time. My suspicion is it will -- AxisKeys works by turning every array into a functor that returns the output associated with an index, so maybe it compiles a function for every array?

LoopVectorization may be harder to cut out without losing performance, but I've been considering whether a switch to TensorOperations.jl for calculations might be faster than LoopVectorization+Tullio for larger arrays anyways; unsure how it would affect compile times.

[ParadaCarleton]

julia> @time using AxisKeys
    1.176212 seconds (2.43 M allocations: 158.393 MiB, 1.06% gc time)
julia> @time using DimensionalData
    0.089362 seconds (241.28 k allocations: 17.780 MiB, 25.46% compilation time)

Wow, yeah, the difference is huge. I stuck with AxisKeys because I wanted to have it interface easily with ArviZ.jl, but DimensionalData does exactly the same job in under a tenth of the compile time. Switching should be very easy (if a bit tedious) since the only real difference is DimensionalData uses a slightly different constructor and square brackets for lookup (instead of round brackets).

[ParadaCarleton]

julia> @time using LoopVectorization, Tullio
  1.722977 seconds (4.75 M allocations: 256.594 MiB, 3.16% gc time)
julia> @time using TensorOperations
  0.171592 seconds (425.81 k allocations: 27.868 MiB)

Similar story here. I'll probably switch them out even if the performance drops slightly, although I've been having some very annoying errors -- for some reason trying to replace @tullio with @tensor throws bugs that I can't really make heads or tails of -- things like not understanding := (even though it's definitely included in TensorOps).

[devmotion]

MCMCDiagnosticTools, as we discussed earlier, can be cut out pretty easily. Another way we could cut it out besides what we discussed is just copy/pasting the FFTESS code from MCMCDiagnosticTools so we don't need to take everything else, which might just be the easiest option.

IMO copying code is usually one of the worst alternatives. Regarding FFT, I assume you could/should replace FFTW with AbstractFFTs, as we do in MCMCDiagnosticTools and MCMCChains as well.

[ParadaCarleton]

MCMCDiagnosticTools, as we discussed earlier, can be cut out pretty easily. Another way we could cut it out besides what we discussed is just copy/pasting the FFTESS code from MCMCDiagnosticTools so we don't need to take everything else, which might just be the easiest option.

IMO copying code is usually one of the worst alternatives. Regarding FFT, I assume you could/should replace FFTW with AbstractFFTs, as we do in MCMCDiagnosticTools and MCMCChains as well.

FFT isn't the important bit -- what matters is that we need a function to calculate ESS but don't want to pull everything else in MCMCDiagnosticTools.jl along with it.

Seth suggested letting users calculate the relative effective sample size themselves; I think this is reasonable but it's a bit inconvenient.

[devmotion]

Sure, the FFTW suggestion was just an additional comment.

I haven't benchmarked the precompilation and loading times of MCMCDiagnosticTools but it mainly depends on lightweight packages: https://github.com/devmotion/MCMCDiagnosticTools.jl/blob/main/Project.toml Most of these packages are just interfaces (AbstractFFTs, DataAPI, MLJModelInterface, Tables), stdlibs (LinearAlgebra, Random, Statistics), or standard packages from the math/stats ecosystem that basically every package depends on (SpecialFunctions and StatsBase). The only other dependency is Distributions which is also quite common in stats/probability theory packages. If there are any issues regarding loading/compilation times, it would be good to track them down (invalidations in one of these dependencies or MCMCDiagnosticTools?) and fix them. From just looking at the list of dependencies it is not clear to me why you would want to remove the dependency on MCMCDiagnosticTools and copy some of the code. In my experience usually this leads to diverging code bases, i.e., you will miss out on any improvements or bug fixes in MCMCDiagnosticTools and vice versa; this happened regularly and is still a problem eg with the copied-from-Distributions-and-slightly-modified code in DistributionsAD.

[ParadaCarleton]

Sure, the FFTW suggestion was just an additional comment.

I haven't benchmarked the precompilation and loading times of MCMCDiagnosticTools but it mainly depends on lightweight packages: https://github.com/devmotion/MCMCDiagnosticTools.jl/blob/main/Project.toml Most of these packages are just interfaces (AbstractFFTs, DataAPI, MLJModelInterface, Tables), stdlibs (LinearAlgebra, Random, Statistics), or standard packages from the math/stats ecosystem that basically every package depends on (SpecialFunctions and StatsBase). The only other dependency is Distributions which is also quite common in stats/probability theory packages. If there are any issues regarding loading/compilation times, it would be good to track them down (invalidations in one of these dependencies or MCMCDiagnosticTools?) and fix them. From just looking at the list of dependencies it is not clear to me why you would want to remove the dependency on MCMCDiagnosticTools and copy some of the code. In my experience usually this leads to diverging code bases, i.e., you will miss out on any improvements or bug fixes in MCMCDiagnosticTools and vice versa; this happened regularly and is still a problem eg with the copied-from-Distributions-and-slightly-modified code in DistributionsAD.

Makes sense. From what I can see, MCMCDiagnosticTools' load time comes primarily from Distributions.jl. Do you happen to know where Distributions.jl is used in MCMCDiagnosticTools? Maybe if it's a few lines that need, say, a Normal PDF calculation, that could just be hand-coded instead to remove the dependency.

Another way to cut load times might be switching from FFTESS to another ESS method so we can drop FFTW, but I have no idea how they compare performance/accuracy-wise. (Using AbstractFFTs doesn't make much difference since users will still have to load their own FFT library.)

[devmotion]

Another way to cut load times might be switching from FFTESS to another ESS method so we can drop FFTW, but I have no idea how they compare performance/accuracy-wise.

Performance-wise in most of my benchmarks computing the ESS with FFT is the worst out of all methods.

(Using AbstractFFTs doesn't make much difference since users will still have to load their own FFT library.)

That's the main point of using AbstractFFTs - users can pick whatever FFT library they want that implements the interface.

[devmotion]

Some benchmarks: TuringLang/MCMCChains.jl#156 and TuringLang/MCMCChains.jl#215 FFT (using FFTW) was by far the slowest method.

[devmotion]

This is also why it's neither the default nor recommended 😄

[ParadaCarleton]

This is also why it's neither the default nor recommended

Haha, makes sense! I just copied it from R's loo package, since I figured they must have some good reason for using it. I'll swap it out for another one then -- is the default the best to stick with?

[devmotion]

According to the linked benchmarks, yes.

I checked where exactly Distributions is used in MCMCDiagnosticTools. It turns out mainly in the discrete diagnostics (for sampling and for evaluating CDFs, the former could maybe be replaced by some, possibly less optimized, calls to StatsBase and the latter by calling into StatsFuns) but unfortunately it is crucial for the Rstar diagnostic: for probabilistic predictions we return a scaled Poisson Binomial distribution (https://github.com/devmotion/MCMCDiagnosticTools.jl/blob/90f502dec5f116f7c12a42df1aa285137dffe6af/src/rstar.jl#L129). So I don't think it can be removed as a dependency.

BTW the dependencies of Distributions are lightweight as well (https://github.com/JuliaStats/Distributions.jl/blob/master/Project.toml), maybe apart from StatsFuns which still depends on Rmath (there are JuliaStats/StatsFuns.jl#125 and JuliaStats/StatsFuns.jl#113 though), so probably any performance issues are caused by Distributions itself. Since it is loaded by most stats/probability theory/Turing packages anyway I wonder how helpful it is to try to remove it though.