Add meta-information to measurable variables

[ricardoV94]

Some of the operations done on measurable variables have implications that are important for other intermediate rewrites. Two come to mind:

1. Dim-shuffling alters the support dimension of Measurable Variables, and because of this we can only safely measure dimshuffles of pure RandomVariables (as they have a "contract" that support dimensions are always to the right): Only introduce MeasurableDimShuffles when input is a RandomVariable #164

2. Some operations convert a continuous variable to a "discrete variable" (whose output may however be a float). This happens with censoring (clipping or rounding) where the output of these operations can/must be of float type, but the variable itself is no longer (strictly) continuous. Other rewrites like "find_measurable_transforms" currently only produce correct IR for continuous variables, and they try to naively reject discrete variables by checking the dtype of the input variables:

   aeppl/aeppl/transforms.py

   Lines 289 to 291 in 46f6775

   	# Do not apply rewrite to discrete variables
   	if measurable_input.type.dtype.startswith("int"):
   	return None

My suggestion is to add meta-information to our intermediate Measurable Variables, namely whether they are continuous/discrete/mixed and the position of their support dimensions. This should be pretty easy as in most IR we can just say they are the same as those from the measurable inputs.

[brandonwillard]

This sounds related to what we effectively tried to do in the beginning: make things fit into the RandomVariable interface.

[ricardoV94]

Do you see a problem with it?

[brandonwillard]

Not necessarily. If anything, it seems like it could lead right back to the interface parts of aesara-devs/aesara#695 (i.e. the signature elements). We definitely don't want yet another piecemeal Op interface (e.g. like nfunc_spec and/or RandomVariable.[ndim_supp, ndims_params]), so, if—for instance—signature information is a part of this, then we need to work on or in the same direction as those issues first.

[rlouf]

I see three benefits to using the existing MeasurableVariable or a new ProbabilityMeasure object to represent measures in AePPL:

• We are not tied to the limitation that Aesara must follow NumPy's API. We could add extra distributions with no impact upstream; conversely there wouldn't be the temptation to do more than NumPy in Aesara. This removes pressure on both libraries;
• We can tie information to these objects that are necessary in AePPL's context without affecting RandomVariables. We can attach densities directly to measures, and add type information: type (support) and parameter constraints.
• We pave the way for a more consistent type system; RandomVariables are not random variables in the measure theory sense, but merely operators that return a random number given a rng state and a set of parameters.

This object of type $\mathbb{M}\alpha$ could subclass RandomVariable (but we don't have to), and we would probably need to subclass RandomStream as well to keep the same interface. We can convert Aesara's RandomVariables to probability measures when building the IR so users can build their models with Aesara.

To address Ricardo's second point we could implement objects that represent transformations between measures, which can be of different types $\mathbb{M}\alpha \to \mathbb{M}\beta$ and lift the corresponding Aesara operators $\alpha \to \beta$ when building the IR. Users would have a choice between working in AePPL's IR directly by defining the measures and calling transformations between measures, or work with Aesara's operators.

Besides the amount of work that this represents I don't see any downside to this approach.

Edit: I don't have a clear understanding of how sampling and density relate mathematically, and we may want to keep RandomVariable and ProbabilityMeasure distinct types with different functions. With ValuedVariable keeping track of the relation during disintegration.