Allow step methods to work with variables that have dynamic sizes

[ricardoV94]

#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]

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

[twiecki]

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

[hvasbath]

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]

+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]

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.