TuringLang / TuringLang/DynamicPPL.jl
VNT: VarInfo mop-up
Nobody has claimed this yet.
- 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!!andlink!!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 improvedextract_priorsthatlink!!uses, but they weren't the culprit. I think this warrants more careful profiling. (Edit: Both functions are significantly faster, often by over 10x.) -
mergecould be optimised to not revert toLinked=nothingso easily. This probably warrants writing a function to check the link status of all variables and use that to setLinked. (A bit likeconcretise_eltypein a way.) - Try removing all or some of the
@inboundsannotations and see how they affect performance. See https://github.com/TuringLang/DynamicPPL.jl/pull/1183#issuecomment-3755257644. -
keyscould be made it faster. Not a priority. (Edit:keysis 2-10x faster than when this issue was written.) -
subsetcould 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:subsetis 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 removingvalues_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
transformin VarInfo. We almost always pull the transform from the distribution intilde_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 ofTransformedValueanymore. - The above will also allow to get rid of some dangerous functions like
unflatten!!andvalues_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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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