Refactor old Distribution base class

[ricardoV94]

PyMC distribution classes are weird objects that hold RandomVariables, logp, logcdf and moment methods together (basically doing runtime dispatching) and manage most of the non-RandomVariable kwargs that users are familiar with (observed, transformed, size/dims) and behind the scenes actions like registration in the model.

This exists mostly for backwards compatibility with V3 and ease of developer refactoring, but the current result is far from pretty.

We need to figure out a more elegant/permanent architecture now that many things that existed to accommodate V3 limitations no longer hold.

Distribution

Distribution is currently performing the following tasks:

pymc/pymc/distributions/distribution.py

Line 135 in 75ea2a8

class Distribution(metaclass=DistributionMeta):

1. Input validation:
   1. Raising FutureWarnings for testval kwarg
   2. Raising TypeError when distribution is initialized outside of a Model context
   3. Raising TypeError when name is not given to a distribution
   4. Raising ValueError when more than one of dims/shape/size is given
2. Convert alternative parametrizations to standard parametrization (e.g, tau -> sigma). This is done by the .dist methods.
3. Add informative attribute errors for deprecated logp, logcdf, random methods
4. Resize the final RV based on observed, shape, dims or size
5. Provides the .dist() API to create an unnamed RV that is not registered in the model. This type of variables is necessary for use in Potentials and other distribution factories that use RVs as building blocks such as Bound and Censored distributions, as well as Mixtures and Timeseries once they get refactored for V4

DistributionMeta

In addition we have a DistributionMeta that does the following:

pymc/pymc/distributions/distribution.py

Line 70 in 75ea2a8

class DistributionMeta(ABCMeta):

1. Dispatch the logp, logcdf, moment, default_transform methods defined in the old PyMC distributions to apply to the respective rv_op
2. Register the rv_op type as subclass of the old style PyMC distribution, so that V3 Discrete/Continuous subclass checks still work?

isinstance(pm.Normal.dist().owner.op, pm.distributions.Continuous)  # True

If we want to get rid of Distribution we probably need to statically dispatch our methods to the respective rv_op. That is nothing special, and is how we do it for aeppl from the get go: https://github.com/aesara-devs/aeppl/blob/38d0c2ea4ecf8505f85317047089ab9999d2f78e/aeppl/logprob.py#L104-L130

[ricardoV94]

They serve the same goal the "standard" Distribution class is doing right now: basically manage all non distribution kwargs: shape/dims, transformed, observed; and behind the scenes magic things like value variables and rngs

The base Distribution class was only ever supposed to make the hand-off to the Model class, which would then handle all the generic tasks (e.g. value variables, RNGs, shape/dims, etc.), and the rest was just for backward compatibility (e.g. the now completely unnecessary Disribution.dist interface). Notice how instances of these classes are never made, and that only static/class/type-level functions are used.

By extending Distribution and adding more logic to those classes, it could become considerably more difficult to unwind those temporary backward compatibility-only choices, and they'll eventually become a permanent and confusing part of v4's design.

This isn't the time or place to start addressing all that, but these are the kinds of design considerations that need to go alongside changes/additions to the relevant areas of code (i.e. Distribution-related code).

Originally posted by @brandonwillard in #5169 (comment)

[fonnesbeck]

Do we have a spec in mind for a replacement class? Maybe an explicit list of what needs to be removed and added with respect to the current class would be a good place to start.

[ricardoV94]

Do we have a spec in mind for a replacement class? Maybe an explicit list of what needs to be removed and added with respect to the current class would be a good place to start.

I will update the issue tomorrow with those

[canyon289]

to accommodate V3 limitations no longer hold.
If youre easily able to list these no longer valid limitations while writing the doc I would also appreciate that (to help with my personal understanding). No obligation if itll take you a lot of time

[ricardoV94]

Updated the top post to mention every function of Distribution classes that I could find. Also highlighted some of the points I think will be a bit more challenging, or may need more drastic API changes.

[fonnesbeck]

eval() should be aliased to something more intuitive like sample() or draw() where appropriate.

[ricardoV94]

eval() should be aliased to something more intuitive like sample() or draw() where appropriate.

eval is just the standard Aesara debug feature that exists for any node. We can have a function wrapper that compiles the "proper" aesara function and takes a given number of draws. For analogy with V3 it could be named pm.random, but I like the draw name better.

The question of RNGs with default updates exists anyway, but could arguably be offloaded to pymc.aesaraf.compile_pymc which already does this for Simulator variables anyway, and which pm.random/ pm.draw would call. That would be one less thing that Distribution has to be concerned with.

pymc/pymc/aesaraf.py

Line 968 in 75ea2a8

# Set the default update of a NoDistribution RNG so that it is automatically

[fonnesbeck]

Yeah, I like a functional approach there, and I agree draw is better than sample since the latter is already associated with MCMC sampling.

[michaelosthege]

What is the selling point of it?
It's a big big refactoring that forces us into an API breakage.
And after adding all those decorators we'd probably end up with more lines of code per distribution; considerably harder to comprehend.

---

First of all, if we create no instances of any PyMC distribution. Why do they exist at all?

We distinguish between variables that are

• registered by name in a model: pm.MyDistribution("name", some, params) or
• created as unregistered, auxiliary variables, that may still contribute to a model graph: pm.MyDistribution.dist(some params)

Pure Aesara RV Ops don't make this distinction (point 7).

Then PyMC distributions have a much richer API & behavior compared to Aesara RandomVariable Ops (points 1,2,3,5).

So we need to wrap the Op.
But should these adapters be a class MyDistribution(...) or a def MyDistribution(...)?

Neither MyDistribution nor MyDistribution.dist() seem to get much love, and generally it feels like there's a recent hype for functional API designs.
But class vs. function aside, let's consider just the resulting user API syntax:

1. pm.Normal("name", 0, 1) and pm.Normal.dist(0, 1) → the current API. Done with class Normal(pm.Distribution) because otherwise we'd be monkey-patching things onto functions and nobody wants that. We can still rename .dist() to .cool() or something.
2. pm.Normal("name", 0, 1) and pm.Normal(0, 1) → doesn't work as long as name, mu, sd etc. are positional. Same for having a kwarg like pm.Normal(0, 1, register=False). In C# this would be trivial, by the way.
3. pm.Normal("name", 0, 1) and pm.normal(0, 1) → Both could be functions, but maybe also easy to confuse?

So we could move away from class MyDistribution(Distribution) only if we switch to syntax 3.

Then every distribution needs to get a def mydistribution(...) and def MyDistribution(...).
The latter could be a one-liner: MyDistribution = pm.as_rich_RV(mydistribution).

Then we'd have to dispatch logp and get_moment etc. onto the RV Op as @ricardoV94 suggested.
The resulting code would no longer be grouped by distribution, and have a ton of dispatch decorators doing the job that's currently done by DistributionMeta.

---

6. Taking care of default updates for the RandomState variables so that returned variables "look" random by default.

This once bit me really hard when I tried to demo Aesara to someone.
I wanted to show how conveniently one can do some symbolic calculations and .eval() things as random numbers.
...took me an hour to understand that I was use to a really awesome behavior that was actually a PyMC, not an Aesara feature.

If Aesara made the non-deterministic .eval() the default---like with PyMC RVs---it would be a lot more beginner friendly & interesting for people doing, for example Monte Carlo sampling.

[ricardoV94]

pm.Normal("name", 0, 1) and pm.Normal(0, 1) → doesn't work as long as name, mu, sd etc. are positional. Same for having a kwarg like pm.Normal(0, 1, register=False). In C# this would be trivial, by the way.

Can't we just check if the first variable is a string and react accordingly?