Batched MvNormal distribution

[Sayam753]

Thank your for opening a PR!

Initializing the work on #5383

Blocked by aesara-devs/aesara#798
EDIT: Now, I am correctly using at.broadcast_shape function.

Depending on what your PR does, here are a few things you might want to address in the description:

• what are the (breaking) changes that this PR makes?
  This PR generalizes the MvNormal distribution.

The parameters mu and cov were restricted to 2d. To remove this restriction, I created a custom Op that operates on batches of square matrices to compute their inverses in vectorized fashion.

With the changes in this PR, I think that random variable object for MvNormal distribution can be created and random samples can be drawn. The logp computations still need to be sorted out. So, this PR is WIP.

• important background, or details about the implementation
• are the changes—especially new features—covered by tests and docstrings?
  Will add them soon.
• linting/style checks have been run
• consider adding/updating relevant example notebooks
• right before it's ready to merge, mention the PR in the RELEASE-NOTES.md

[Sayam753]

Two questions:

• Is there an existing function in aesara that can inverse batches of square matrices in a vectorized way?
• If no, I have added this Op in a new file. Where should this class ideally be placed?

[ricardoV94]

distributions/dist_math.py?

[ricardoV94]

If numpy has support for batched matrix inversion but not Aesara, we can open an issue there. Here is a similar case: aesara-devs/aesara#791

[ricardoV94]

Actually the Aesara Op is just doing the same as this one but has an unnecessary ndim check?

We should just open a PR to fix it there

[codecov]

Codecov Report

Merging #5424 (5ee641f) into main (73add2d) will decrease coverage by 6.64%.
The diff coverage is 68.88%.

@@            Coverage Diff             @@
##             main    #5424      +/-   ##
==========================================
- Coverage   81.38%   74.74%   -6.65%     
==========================================
  Files          82       83       +1     
  Lines       14207    14241      +34     
==========================================
- Hits        11563    10645     -918     
- Misses       2644     3596     +952     

Impacted Files	Coverage Δ
pymc/distributions/multivariate_utils.py	61.76% <61.76%> (ø)
pymc/distributions/multivariate.py	71.76% <90.90%> (-19.77%)	⬇️
pymc/backends/tracetab.py	28.57% <0.00%> (-71.43%)	⬇️
pymc/func_utils.py	22.50% <0.00%> (-65.00%)	⬇️
pymc/printing.py	22.22% <0.00%> (-63.64%)	⬇️
pymc/distributions/censored.py	37.28% <0.00%> (-55.94%)	⬇️
pymc/distributions/bound.py	44.32% <0.00%> (-55.68%)	⬇️
pymc/exceptions.py	58.62% <0.00%> (-41.38%)	⬇️
pymc/distributions/discrete.py	68.64% <0.00%> (-31.12%)	⬇️
pymc/distributions/continuous.py	75.88% <0.00%> (-21.05%)	⬇️
... and 13 more

[ricardoV94]

@Sayam753 @purna135 Shall we close this one while the required changes in Aesara are not implemented?

[Sayam753]

Thanks for the ping Ricardo.
Yes. This can be closed now and revived when we have Blockwise implemented in Aesara.