TuringLang / TuringLang/DynamicPPL.jl

VNT: VarInfo mop-up

Open
#1,201 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

data-structures enhancement
Dominant language
Julia
Stars
286
Forks
41
Avg merge
1d 11h
Merged PRs (30d)
34

Description

Due to time constraints on my part, and to avoid ballooning the PR, we are merging #1183 into its base branch (though not breaking) even though some issues remain. Listing those issues here:

  • unflatten!! and link!! remain surprisingly slow for models that have IndexLenses. I spent an hour on this but failed to understand why. I tried improving type stability within those functions by getting rid of some closures, and I also improved extract_priors that link!! uses, but they weren't the culprit. I think this warrants more careful profiling. (Edit: Both functions are significantly faster, often by over 10x.)
  • merge could be optimised to not revert to Linked=nothing so easily. This probably warrants writing a function to check the link status of all variables and use that to set Linked. (A bit like concretise_eltype in a way.)
  • Try removing all or some of the @inbounds annotations and see how they affect performance. See https://github.com/TuringLang/DynamicPPL.jl/pull/1183#issuecomment-3755257644.
  • keys could be made it faster. Not a priority. (Edit: keys is 2-10x faster than when this issue was written.)
  • subset could be made faster, see the comment in it. It's already faster than the old version in most cases though, so not a priority. (Edit: subset is 1-2x faster than when this issue was written.)
  • See the benchmarking script I wrote and see if some of it could feed into our benchmark suite.
  • Check that AD performance is fine. I've been benchmarking primal evals so far. I don't expect anything to go wrong there, but should do a quick check.
  • Reimplement Base.rand(::Model) using a VAIMAcc, and make it return a VNT. That will allow removing values_as(::VarInfo, ::NamedTuple), which shouldn't exist.
  • Once the performance optimisations are all done, add a note in HISTORY.md for what's changed performance-wise. There's a placeholder in there already.

There are also bigger issues that might not make it to this release, and probably deserve their own issues:

  • Stop vectorising all values in VarInfo, now that the backend storage no longer requires it.
  • Stop storing transform in VarInfo. We almost always pull the transform from the distribution in tilde_assume!! anyway, because the stored transform may be out of date. Safer to not even store it, and makes type stability easier because linking doesn't have to change the type of TransformedValue anymore.
  • The above will also allow to get rid of some dangerous functions like unflatten!! and values_as(::VarInfo, ::Colon).
  • Start using shadow arrays, see #1194
  • Reimplement StaticTransformation. It was a bit broken/risky, and now it's broken/risky in a slightly different way. See comments in the function.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the benchmarking script linked from PR #1183 and inspect the current VarInfo implementation and benchmark suite. Determine which benchmark cases can be incorporated, then assess the unchecked vectorisation and cleanup items. Done means the selected benchmarks are integrated and any remaining scope is resolved or split into focused issues.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, performance
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.