Skip to content

VNT: VarInfo mop-up #1201

@mhauru

Description

@mhauru

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.
  • 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 VNT Part 5: VarNamedTuple as VarInfo #1183 (comment).
  • keys could be made it faster. Not a priority.
  • 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.
  • 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 VNT: Handling arrays generally #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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions