Increase support for batched multivariate distributions

[ricardoV94]

The logp of several multivariate distributions does not work (or is not tested) for arbitrarily batched dimensions. Some cases I could confirm include:

• MvNormal - Allow batched parameters in MvNormal and MvStudentT distributions #6897
• MvStudentT Allow batched parameters in MvNormal and MvStudentT distributions #6897
• KroneckerNormal
• MatrixNormal
• LKJCorr
• LKJCholeskyCov
• StickBreakingWeights (alpha) Allow for batched alpha in StickBreakingWeights #6042
• Wishart (may be fine to ignore)
• CARRV (maybe better to leave this one as is)

Reproducible code:

size = (4, 3)
pm.logp(pm.MvNormal.dist(mu=np.ones(2), cov=np.eye(2), size=size), np.ones((*size, 2)))
# ValueError: Invalid dimension for value: 3
pm.logp(pm.MvStudentT.dist(nu=3, mu=np.ones(2), cov=np.eye(2), size=size), np.ones((*size, 2)))
# ValueError: Invalid dimension for value: 3

Distributions that already support (and have tests for) arbitrary shapes

• Dirichlet
• Multinomial
• DirichletMultinomial - see Tweak DirichletMultinomial logp and refactor some multivariate logp tests #5234

[purna135]

I'd like to work on this, could you please point me in the right direction?

[Sayam753]

I think one can start by exploring the implementation of Random Variables in v4. @ricardoV94 did a great job of explaining design steps here.

• Then, gather all Ops from Aesara that are used in the implementations of multivariate distributions.
• Generalize the Ops to work beyond 2d inputs.
• Make sure the updated Ops work readily / Make use of Aesara's broadcasting capabilities to integrate the updated Ops in multivariate distributions.
• Update any/all test cases related to these Ops and multivariate distributions.
• Bonus here would be to see how mixture distributions will play with generalized multivariate distributions.

[ricardoV94]

* Bonus here would be to see how mixture distributions will play with generalized multivariate distributions.

They should work fine out of the box. We have tests for batched Dirichlet for instance. Only limitation is whether the distribution accepts batched values in the logp

[Sayam753]

For mixtures, I think that generalization of multivariate distributions is a major improvement as compared to v3. For creating a mixture of lets say multivariate normals, there needs to be a loop in v3 to create mixture components. On the other hand, I can just specify the mixture axis in v4 if multivariate distributions support working with batch data.

[ricardoV94]

For mixtures, I think that generalization of multivariate distributions is a major improvement as compared to v3. For creating a mixture of lets say multivariate normals, there needs to be a loop in v3 to create mixture components. On the other hand, I can just specify the mixture axis in v4 if multivariate distributions support working with batch data.

Yeah, I was just saying that multivariate mixtures already work in V4, nothing else should be needed. Am I missing something?

[purna135]

Thank you so much, @Sayam753 and @ricardoV94; I will follow these steps and let you know if any assistance is required : )

[purna135]

• StickBreakingWeights (alpha, but not k, could be batched)

Hi @ricardoV94. StickBreakingWeights can now take batched alpha(example here) after removing some simple checks for alpha. Should we also generalize K for StickBreakingWeights ? Is there a particular instance of K for which we are not generalizing?

[ricardoV94]

We can't generalize k because that would mean draws with different sizes and Aesara cannot handle ragged (non rectangular) tensors.