Implement several RandomVariables as SymbolicRandomVariables

[ricardoV94]

Description

Implement several RandomVariables as SymbolicRandomVariables

This allows sampling from multiple backends without having to dispatch for each one

Also:

• Allow signature to handle rng and size arguments explicitly.
• Move rv_op method to the SymbolicRandomVariable class and get rid of dummy inputs logic (it was needed in previous versions of PyTensor)
• Allow dispatch methods without filtering of inputs for SymbolicRandomVariable distributions

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7239.org.readthedocs.build/en/7239/

[jessegrabowski]

Seems like this is mostly shuffling around and simplifying stuff that already exists, so if tests pass it looks great. I'm a bit unclear on how the signatures for RVs work, was that changed in this PR or another one? The square brackets in particular throw me off. There's also some opportunities for refactoring duplicated code.

[ricardoV94]

Seems like this is mostly shuffling around and simplifying stuff that already exists, so if tests pass it looks great.

There are two commits in this PR. The first one shuffles and refactors things so that the second commit (converting existing RandomVariables to SymbolicRandomVariables) is done mostly without hassle.

I'm a bit unclear on how the signatures for RVs work, was that changed in this PR or another one? The square brackets in particular throw me off.

I think this is the kind of signature we should use in PyTensor, so I was testing it here. The numpy vectorize signature doesn't really handle stuff like rng, size, axis, so I think we should take the initative here. In a previous PR I pretend those were scalars, so stuff looked like (),(),(),()->(),(), which is both fake and I think less readable than [rng],[size],(),()->[rng],(). Wdyt? This would ultimately replace the ndim_supp and ndims_params which is very strict and only works because we assume all RVs have the same non-numerical signature. This is not useful for SymbolicRandomVariables that can have multiple (or zero) RNGs and only optionally size.

There's also some opportunities for refactoring duplicated code.

I'm weary of some helpers, but I can optimize for a bit less DRY

[ricardoV94]

@jessegrabowski I simplified the logic needed to build standard SymbolicRandomVariables, what do you think?

[ricardoV94]

Some bit of blackmagic to allow dynamic properties from the signature, to work both when signatures are class attributes or only instance attributes. Let me know if you see a better/ less magical solution.

[jessegrabowski]

Why is it needed though?

[ricardoV94]

Sometimes we have the signature fixed in the class (DM), sometimes only known at runtime (Mixture). This allows us to parse ndims_params, ndim_supp, default_outputs from the signature regardless of when it's defined.

[jessegrabowski]

Would it be better to have a class_signature (always defined) and a instance_signature (default None), and drop back to the class_signature if instance_signature is None?

[ricardoV94]

A class signature cannot always be defined. The point is I want to extract ndim_supp from either a class signature (without having to initiate said class) or, if it cannot be known in advance, from an instance signature. I need a static class method or instance method that can read either the class or instance signature, this achieves exactly that.

[ricardoV94]

Do I really need this? Not so much. Alternatives are:

1. Hard code ndim_supp/ndims_params/default_output when they can be known in advance. It's just redundant/error prone because all info exists in the signature.
2. Don't provide any of these until the Op is defined. Drawback? I can't use these lines as straightfoward:

We have to create dummy Ops with empty size just just to find the ndim_supp:

pymc/pymc/distributions/distribution.py

Lines 486 to 489 in 5f95fc2

	# SymbolicRVs don't have `ndim_supp` until they are created
	ndim_supp = getattr(cls.rv_op, "ndim_supp", None)
	if ndim_supp is None:
	ndim_supp = cls.rv_op(*dist_params, **kwargs).owner.op.ndim_supp

I can't write classmethods of the new SymbolicRandomVariables like this:

        if rv_size_is_none(size):
            size = implicit_size_from_params(b, kappa, mu, ndims_params=cls.ndims_params)

Because cls.ndims_params doesn't exist until we create the Op below.

I could obviously parse ndims_params/ndim_supp on the spot from the signature. It's just a small convenience. But OTOH it's also just 5 lines of Python

[ricardoV94]

@jessegrabowski My last commit reverts the magic solution for the cleanest alternative I can think of. Which one do you prefer?

[ricardoV94]

I removed the non-magic commit, for reference it's here: 20e6e3b