Refactoring _print_name for certain RVs and specifying rv_types in their distributions

[larryshamalama]

What is this PR about?
This PR creates "new" PyMC RandomVariables for those that have an abbreviate _print_name such that we don't have to rely on the RV's class name for Latex and Graphviz displays.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• ...

Bugfixes / New features

• ...

Docs / Maintenance

Subclass several random variables to have a more readable _print_name.

Closes #6201

[codecov]

Codecov Report

Merging #6219 (1205bd8) into main (7bad057) will decrease coverage by 0.00%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6219      +/-   ##
==========================================
- Coverage   93.78%   93.78%   -0.01%     
==========================================
  Files         101      101              
  Lines       22187    22235      +48     
==========================================
+ Hits        20809    20854      +45     
- Misses       1378     1381       +3     

Impacted Files	Coverage Δ
pymc/distributions/continuous.py	97.56% <100.00%> (+0.06%)	⬆️
pymc/distributions/discrete.py	99.25% <100.00%> (+0.03%)	⬆️
pymc/distributions/distribution.py	95.00% <100.00%> (ø)
pymc/distributions/multivariate.py	92.32% <100.00%> (+0.06%)	⬆️
pymc/tests/distributions/test_logprob.py	100.00% <100.00%> (ø)
pymc/parallel_sampling.py	85.52% <0.00%> (-0.99%)	⬇️

[twiecki]

Can't we just set _print_name on the imported aesara rvs?

[ricardoV94]

Can't we just set _print_name on the imported aesara rvs?

That would have side-effects on the Aesara side (if it even works)

[larryshamalama]

Letting tests run again. Hopefully nothing fails.

Are there any additional tests to add in light of these new Ops?

[ricardoV94]

I am not super happy with this (even though I suggested it). It's a bit worrisome that our dispatched methods will work for the subclasses but not the original classes, when all we are doing is changing a name.

We could do something more clever, where we still dispatch on the Aesara classes but return the PyMC ones from Distribution when those are specified. Then both objects will have logp/logcdf/moments/ etc...

That would have avoided the concerns about the tests where we create new classes.

We could use the rv_type property, and dispatch on that. So Normal would be

class Normal(Continuous):
  rv_op = pymc_normal
  rv_type = NormalRV  # Aesara
  ...

And then tweak the Metaclass of Distribution to use rv_type for dispatching when provided explicitly, and otherwise default to the old behavior of obtaining the rv_type from rv_op. WDYT?

[larryshamalama]

I am not super happy with this (even though I suggested it). It's a bit worrisome that our dispatched methods will work for the subclasses but not the original classes, when all we are doing is changing a name.

I understand. I also had some doubts, but at least we have a clearer picture of how the changes in this PR would affect the codebase.

We could do something more clever, where we still dispatch on the Aesara classes but return the PyMC ones from Distribution when those are specified. Then both objects will have logp/logcdf/moments/ etc...

That would have avoided the concerns about the tests where we create new classes.

We could use the rv_type property, and dispatch on that. So Normal would be

class Normal(Continuous):
  rv_op = pymc_normal
  rv_type = NormalRV  # Aesara
  ...

And then tweak the Metaclass of Distribution to use rv_type for dispatching when provided explicitly, and otherwise default to the old behavior of obtaining the rv_type from rv_op. WDYT?

On top of my head, both solutions would work, but I'm not sure how I would dispatch the Aesara classes given rv_op to be the newly created ones (like PyMCNormalRV).

pymc/pymc/distributions/distribution.py

Lines 107 to 114 in 13dfeb2

	if isinstance(rv_op, RandomVariable):
	rv_type = type(rv_op)
	clsdict["rv_type"] = rv_type
	new_cls = super().__new__(cls, name, bases, clsdict)
	if rv_type is not None:
	# Create dispatch functions

I believe that the dispatch is being done in these lines and perhaps there is a more clever way than checking if __class__.__name__.startswith("PyMC")... Is this what you meant?

Again, without trying much, I feel like using rv_type would be an easier solution. However, would this have any spillover complications in other areas of the codebase?

[ricardoV94]

@larryshamalama what I meant was something like:

        rv_type = clsdict.setdefault("rv_type", None)

        if rv_type is None and isinstance(rv_op, RandomVariable):
            rv_type = type(rv_op)
            clsdict["rv_type"] = rv_type

Then when we set:

class Normal(Continuous):
  rv_op = pymc_normal
  rv_type = NormalRV  # Aesara

pm.Normal will still return a PyMCNormalRV, but all the dispatching will be done on the original baseclass. Dispatching works for subclasses but not parent classes (which makes sense).

[larryshamalama]

Okay, I can get to it

[ricardoV94]

@larryshamalama Can you include the default changelist from the PR template? We use those for the auto-generated release notes

[larryshamalama]

Done! Thanks for the suggestion, I will start doing this more often

[ricardoV94]

Thanks @larryshamalama, now the print_name_ for SymbolicRVs should be more useful