Can't sample a model after resizing an RV

[michaelosthege]

Description

First of all this is not a regression w.r.t. PyMC3 v3, since model resizing did not even work to begin with.

But with the new RandomVariable based model graphs, the size of an RV is no longer cut in stone and we're promoting that as a new feature, so users might try it and run into this issue.

I'm opening this so it's documented.

Here's a minimal example:

import pymc3 as pm

def run():
    with pm.Model() as pmodel:
        data = pm.Data("data", [1, 2, 3])
        x = pm.Normal("x", size=data.shape[0])
        pm.Normal("likelihood", mu=x, observed=data)

        sample_kwargs = dict(cores=1, chains=1, tune=2, draws=2)
        pm.sample(**sample_kwargs, step=pm.Metropolis())

        # resize the model
        data.set_value([1, 2, 3, 4])
        pm.sample(**sample_kwargs, step=pm.Metropolis())

if __name__ == "__main__":
    run()

Traceback

ValueError: Input dimension mismatch. (input[1].shape[0] = 4, input[2].shape[0] = 3)
Apply node that caused the error: Elemwise{Composite{(i0 + (-sqr((Cast{float64}(i1) - i2))))}}(TensorConstant{(1,) of -1..0664093453}, data, x)
Toposort index: 0
Inputs types: [TensorType(float64, (True,)), TensorType(int32, vector), TensorType(float64, vector)]
Inputs shapes: [(1,), (4,), (3,)]
Inputs strides: [(8,), (4,), (8,)]
Inputs values: [array([-1.83787707]), array([1, 2, 3, 4]), array([-1.2947119 ,  0.10578338, -0.59358436])]
Outputs clients: [[Sum{acc_dtype=float64}(Elemwise{Composite{(i0 + (-sqr((Cast{float64}(i1) - i2))))}}.0)]]

HINT: Re-running with most Aesara optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the Aesara flag 'optimizer=fast_compile'. If that does not work, Aesara optimizations can be disabled with 'optimizer=None'.
HINT: Use the Aesara flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

Why this happens

The initial values are eagerly set/drawn/evaluated to numeric ndarrays and stored in the model.initial_values dictionary.
If the corresponding RV changes its size later on, the initial value is not updated.

Potential solutions

We could prefer to keep track of symbolic initial values (or None) and only evaluate/draw them at the start of sampling.

User-provided numeric initial values would still restrict the RV to its initial size, of course, but it should fix most use cases.

Versions and main components

• PyMC3 Version: main (also on 725d798, so no it is not because of latest initval changes)

[brandonwillard]

If the random variables in a model are changed, then their corresponding initial values need to be updated. You can use Model.set_initval for that.

Otherwise, you're reusing a step instance; one that has state that is no longer valid, just like the updated variable's initial value. If you create a new pm.Metropolis instance (and update the initial value for the changed variable), it should work as expected.

It would be great if the step methods were stateless, but that's not the case—nor is it a necessary design choice for v4. If you want to open a ticket for that, it would be welcome; otherwise, this issue should be closed. Likewise, if there's an issue with Model.set_initval, please open an issue for that.

[michaelosthege]

The statefulness of step methods has nothing to do with this issue. I changed the example above accordingly.

If you want to open a ticket for that, it would be welcome; otherwise, this issue should be closed.

I don't want to open a ticket for that right now and this issue has nothing to do with it, so there's no reason to close it because of an unrelated quirk.

If the random variables in a model are changed, then their corresponding initial values need to be updated. You can use Model.set_initval for that.

Sure, but in the example above the user did not manually specify initval. We don't advertise Model.set_initval in our quickstart documentation either.
So it's a really bumpy user experience that we can improve, for example, with the solution I proposed above.

[brandonwillard]

The statefulness of step methods has nothing to do with this issue. I changed the example above accordingly.

It has literally everything to do with it, because it prevents PyMC from working even after the initial value problem is fixed.

I don't want to open a ticket for that right now and this issue has nothing to do with it, so there's no reason to close it because of an unrelated quirk.

Can you demonstrate that this example will work without doing anything about the step instance? If not, it definitely has something to do with that.

Sure, but in the example above the user did not manually specify initval. We don't advertise Model.set_initval in our quickstart documentation either.
So it's a really bumpy user experience that we can improve, for example, with the solution I proposed above.

All right, so is this issue about documentation? It seems like it's about an issue you faced while attempting to do something specific, one for which you were provided an explanation and a succinct solution that only involves pre-existing features.

[twiecki]

Otherwise, you're reusing a step instance; one that has state that is no longer valid, just like the updated variable's initial value. If you create a new pm.Metropolis instance (and update the initial value for the changed variable), it should work as expected.

@brandonwillard As it seems to also break with instantiating a new sampler, do you think it's about requiring to set a new initvalue?

Just from an API stand-point I think this use-case would be great to support and not require updating an init-value specifically.

[brandonwillard]

@brandonwillard As it seems to also break with instantiating a new sampler, do you think it's about requiring to set a new initvalue?

No, it's just that those step method instances hold state, and part of that state includes the old size of the updated model variable. In other words, it's the design of the step method instances that's the problem. If they were made stateless in the relevant way, this wouldn't be an issue.

Just from an API stand-point I think this use-case would be great to support and not require updating an init-value specifically.

No matter what, the initial values need to be updated when a model variable's shape is changed. That's the issue, and there's no avoiding it.

If we want to automate that—when possible—that's fine. For instance, pm.Data could attempt to trigger such an action when set_value is called.

Such automations will end up being rather context dependent, because, for example, adding that automation to pm.Data.set_value wouldn't cover the more general case of an aesara.shared variable. Anything that attempts to do this more generally will most likely involve a much larger set of changes with questionable utility-to-effort ratios.

We should focus on making it easy to update/set initial values, and making people aware of the need to do that. When the former is covered, we can also make it easier to implement some of these automations using the same machinery.

[twiecki]

No, it's just that those step method instances hold state, and part of that state includes the old size of the updated model variable.

But when we create a new step method (like the updated code above does) we get a new state.

If we want to automate that—when possible—that's fine. For instance, pm.Data could attempt to trigger such an action when set_value is called.

I think that would make sense. 99% of changes to the shape of the model will come through pm.Data. For the other 1%, we can issue an informative error that tells people to update the initval.

[brandonwillard]

But when we create a new step method (like the updated code above does) we get a new state.

Yes, and that should fix that problem and allow the entire example to work after addressing the invalid initial values caused by changing the shape of x (e.g. add a line like pmodel.set_initval(x, ...)).

The things we need to work on could involve, for example, automating downstream init value changes, and those should be possible to do via the existing resources (e.g. Model.set_initval).

[twiecki]

so two things need to happen: 1. update the initval, 2. update the state of the sampler (e.g. by reinstantiating) which is tied to the model shape.

[michaelosthege]

To me it still sounds a lot more straightforward to draw/evaluate/numerify initial values when they are accessed, rather than when they're set.
Then we don't need to implement complex mechanisms for changing (and re-calculating them in topological order) because downstream initial values would be symbolic and resizeable alongside the RVs, no?

I proposed that here and here and in the description of this issue.

I'm sorry, maybe I misunderstood something that was already mentioned in the discussion, or am missing insight on incompatibilities, but so far I am only aware of counterarguments along the lines of:

• needing new functionality / not using existing APIs (pmodel.set_initval)
  ^{that can be used to make the example work if one understands what's going on (bad UX)}
• doesn't work for user-provided numerical initvals
• user could pass Variable initvals that fail to evaluate

whereas on the pro side I see

• resize-compatibility* for the vast majority of initvals (this issue)
  ^{*as long as steppers are re-initialized or properly reset, which we should probably fix regardless}
• moment-based initvals wouldn't need (model aware) evaluation before going into register_rv/set_initval (less complex to fix Reimplement old start value logic #4752)
• also works for user provided symbolic (calculated) initvals without additional overhead (the example from Can't set symbolic initval #4911)
• no need to automate downstream init value changes

Unless someone already knows about a dead end, why not just try and implement it?

[brandonwillard]

To me it still sounds a lot more straightforward to draw/evaluate/numerify initial values when they are accessed, rather than when they're set.

If you're advocating for some form of lazy evaluation, that's not a relevant implementation detail at this stage, nor does it provide any sort of functionality that couldn't be accomplished using a non-lazy approach.

Aside from that, it sounds like you're just saying "recompute initial values when they (might) need to be recomputed", and that's just restating the solution we've been talking about this whole time.

Then we don't need to implement complex mechanisms for changing (and re-calculating them in topological order) because downstream initial values would be symbolic and resizeable alongside the RVs, no?

You're going to need to evaluate graphs in a topologically sound order no matter what, but, more importantly, the approaches you're implicitly comparing are not clear. What exactly is complex and how does whatever you're proposing solve that? Better yet, how is it you can get away with not recalculating initial values in topological order?

The only relevant points made in this issue were

• that initial values need to be recomputed to solve the problem in this issue's example,
• that there already exists at least one means of setting them, and
• that some cases of automatic recalculation are do-able.

If you don't like Model.set_initval, then fine, but are you implying that it's somehow more "complex" than anything else? More importantly, do you know what I was referring to when I mentioned "downstream init value changes"?

I proposed that here and here and in the description of this issue.

These links and this issue's description are broad statements that—at best—imply the production of new initial values when they're needed. Yes, that's the basic idea, but you don't need to rephrase it so that it sounds like you have a particular solution with distinct advantages, or that the problem is actually something else.

Likewise, it doesn't help to have issues that say things like "Can't do X" when "X" actually can be done, but you would like it to be done differently, or be associated with a "proposal" of yours (vague or otherwise).

I'm sorry, maybe I misunderstood something that was already mentioned in the discussion, or am missing insight on incompatibilities, but so far I am only aware of counterarguments along the lines of:

Once again, what exactly are you comparing with those counter-arguments and pros? You've only listed some broad complains, limitations, and features/desires with no concrete connections to any specific proposals. Your mention of pmodel.set_initval implies that you're referring to it as one of the "proposals" you're arguing against, but use of Model.set_initval doesn't preclude any of your pros, making the comparison a simple straw man argument.

Regardless, Model.set_initval isn't particularly important or necessary. We can easily take its current logic and repurpose it to make functions that accept user-specified initial values, compute special values for the rest, and return dicts. This—along with some other more library-wide changes—is a step toward making Model less stateful, but it's not a less "complex mechanism" or more capable approach than using a stateful Model. It's preferable for other reasons, but it's not a distinct solution to anything we've discussed.

Otherwise, it seems like you're forgetting that we're still constrained by v3 requirements and designs, and that your overly general proposals neglect to address any of this. For instance, do you know why the initial values are attached to the Model object? (Hint: it's not because it's the overall best approach.) See my closing comment in the description of #4567 and remember that users need to be able to specify initial values (that was a v3 feature).

[twiecki]

I'm still a bit confused here, but what I gather what we all agree on: A mechanism that resets the initval when pm.Data.set_data() is called would improve UX. So let's just do that, and call Model.set_initval from there.

I'm still, however, not clear on why even this needs to be done. If default initvals were symbolic and tied to the shape of the RV, resizing the RV should automatically resize its initval, no? I clearly don't know the implementation but maybe someone could confirm that "currently default initvals are not symbolic and the reason for that is X."

[ricardoV94]

maybe someone could confirm that "currently default initvals are not symbolic and the reason for that is X."

That's the case indeed. What happens when you add a new RV to a model is that it will compile a function to compute a new initval using any of the already pre-computed initvals of other RVs that are also an input to the new RV (as givens in the Aesara function).

I don't know if there is any clear advantage to this "eager" initval computation as opposed to creating a large function that computes all the initvals at once whenever that's needed.

[twiecki]

Closing in favor of #4924.