Implement MvNormal as cholesky(cov) @ normal

[ricardoV94]

Description

This is much faster, and even more in PyMC models that are usually parametrized with a direct prior on the cholesky.

import pytensor
import pytensor.tensor as pt

srng = pt.random.RandomStream()

x = srng.multivariate_normal([0, 0], [[1, 0.5], [0.5, 1]])
fn = pytensor.function([], x)
%timeit fn()  # 510 µs ± 81.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Decompose cholesky in graph (numpy probably does this under the hood)
A = pt.linalg.cholesky([[1, 0.5], [0.5, 1]])
x = A @ srng.normal(size=(2,))
fn = pytensor.function([], x)
%timeit fn()  # 27.4 µs ± 3.27 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In general we should probably reduce the number of pure RV Ops we have. This allows more optimizations and makes it easier to implement different backends.

We should implement the MvNormal as an OpFromGraph that gets inlined after canonicalization (not as early as the ones with inline=True)

[jessegrabowski]

One negative of using cholesky is that it doesn't allow singular covariance matrices. Any square-root type decomposition will work for generating fast forward samples. For example using SVD:

U, S, V_T = pt.linalg.svd(cov, compute_uv=True)
L = U @ pt.diag(pt.sqrt(S))
x = L @ srng.normal(size=(2,))
fn = pytensor.function([], x)
%timeit fn()  # 8.17 μs ± 76 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

Cholesky is the best in general because it gives access to solve_triangular when we're computing the log-likelihood. I wish we exposed an allow_singular option like scipy/JAX that could let the user accept a performance hit to handle this case.

[ricardoV94]

No reason why we can't have both

[jessegrabowski]

JAX has a method argument that takes SVD, eighsh, or cholesky (default cholesky). I think we should go this route.

[jessegrabowski]

Was this closed by #1203?

[ricardoV94]

No, this is a special case of #1221, to actually represent MvNormal like that, not hidden inside the Op, so PyTensor can reason about that if it wants to