Skip to content

unflatten alone considered harmful #995

@penelopeysm

Description

@penelopeysm

It leaves VarInfo in an inconsistent state

unflatten updates values in the VarInfo, but doesn't update logp to match it.

I think there are vanishingly few use cases where one actually ONLY wants to perform unflatten without re-evaluating the model.

Conversely, there have been a number of bugs arising from the fact that logp isn't updated after unflatten-ing, and the number of such bugs will only increase as time goes on.

One situation where we DID use to call unflatten without evaluating was the general interaction between VarInfos and samplers:

  1. We have a varinfo as the AbstractMCMC state.
  2. We take vi[:] and pass that to the sampler.
  3. The sampler does its thing and returns us values and a log density.
  4. We call unflatten(vi, new_values) and setlogp!!(vi, new_logp). This avoids re-evaluating the model.

Unfortunately, in DynamicPPL 0.37 logp has been separated into prior and likelihood, and there is no (easy) way of apportioning the log density between prior and likelihood without re-evaluating the model.

(Note that I don't consider this a bad change. We have a richer representation of logp and we just need to do more work to make sure that it is correct.)

It's highly dependent on the internal structure of the VarInfo

On top of the danger of leaving the VarInfo in an inconsistent state, the actual effect of unflatten is highly dependent on for example:

  • which variables are present in the model;
  • the order in which they appear;
  • whether they are linked.

Of these points, the third is the most likely to arise in practice (the first and second can be triggered e.g. with dynamically sized models). For example, it means that initialisation can have wrong results if VarInfos are unexpectedly linked, or if the parameters used to initialise them are themselves linked: #793 (comment)

Proposed solution

IMO, all of this suggests that unflatten alone should be considered extremely dangerous. I would not remove it but I would:

  • stick a huge warning on using unflatten alone. I would personally even consider renaming it to something else like unsafe_setvals.
  • create a new method unflatten_and_evaluate which basically does unflatten and evaluate!!, and skips the warning.

For clarity: I am aware that a function like unflatten is necessary. However, that doesn't imply that we should gratuitously use it, nor does it preclude us from creating 'safer' wrappers around it like unflatten_and_evaluate. Pointer arithmetic also exists in Julia and is necessary, however, we certainly don't encourage people to use it whenever they can.

Follow-ups

There are similar issues whereby setindex!! actually sets internal values, getindex always returns user-parametrised values, getindex_internal always returns internal values, BUT getindex(vi, :) returns internal values. #854

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