search

pymc-devs / pymc

Bayesian Modeling and Probabilistic Programming in Python
https://docs.pymc.io/
Other
8.64k stars 1.99k forks source link

Allow step methods to work with variables that have dynamic sizes #5139

Open ricardoV94 opened 2 years ago

ricardoV94 commented 2 years ago

https://github.com/pymc-devs/pymc/issues/4463 raises several points related to this.

Our step methods are built around the assumption that model variables dimensions are fixed during sampling, which we might want to relax. Such refactoring would probably be quite painful, so perhaps we could explore adding some variants of the step methods that allow for this, incrementally. Users would have to manually assign variables they know don't have a fixed shape to them, or we would need some extra logic in the step assignment to check if a variable has non-fixed size.

with pm.Model() as m:
  x = pm.DiscreteUniform('x', 1, 10)
  y = pm.Normal('y', 0, 1, size=x)

  y_step = pm.FANCY_NUTS([y])
  pm.sample(step=y_step)

Some aspects of this might prove trickier. @brandonwillard noted:

pymc3.parallel_sampling._Process

  • The whole parallel sampling approach relies on a single set of fixed-shape shared memory arrays for each random variable, and the processes all appear to write to that set of memory. Unfortunately, that approach doesn't work when the random variables change shape during sampling. This doesn't need to be fixed immediately, but it is an unnecessary restriction caused by this specific approach to multi-processing.
michaelosthege commented 2 years ago

Even before that we need a backend to support this. The NDarray for example does not.

twiecki commented 2 years ago

I'm not sure it's worth it. What type of models require this? Can they even be efficiently sampled?

hvasbath commented 2 years ago

For example in tomographic imaging you need that, to decide whether data allows to resolve if a cell in the image can be discretized finer and thus giving you more information.

drbenvincent commented 1 year ago

+1 for this feature. A specific use case is models whose structure is (partially) inferred. For example, a change-point detection model where the number of change points is inferred. Another would be a Gaussian mixture model where the number of mixtures is to be inferred. That is, assuming my understanding of this issue is correct :)

michaelosthege commented 1 year ago

One step in this direction will be to swap the trace backend to mcbackend such that we can store the dynamically sized draws.

DictToArrayBijection and RaveledVars were not desined to work with dynamic sizes, and refactoring them might be quite difficult. So a step method that can act on dynamic size RVs would probably inherit from BlockedStep directly.