Allow decomposition methods in MvNormal

Author: ricardoV94

pytensor/link/jax/dispatch/random.py

@@ -128,7 +128,6 @@ def jax_sample_fn(op, node):
 @jax_sample_fn.register(ptr.BetaRV)
 @jax_sample_fn.register(ptr.DirichletRV)
 @jax_sample_fn.register(ptr.PoissonRV)
-@jax_sample_fn.register(ptr.MvNormalRV)
 def jax_sample_fn_generic(op, node):
     """Generic JAX implementation of random variables."""
     name = op.name
@@ -173,6 +172,22 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
+@jax_sample_fn.register(ptr.MvNormalRV)
+def jax_sample_mvnormal(op, node):
+    """Generic JAX implementation of random variables."""
+
+    def sample_fn(rng, size, dtype, mean, cov):
+        rng_key = rng["jax_state"]
+        rng_key, sampling_key = jax.random.split(rng_key, 2)
+        sample = jax.random.multivariate_normal(
+            sampling_key, mean, cov, shape=size, dtype=dtype, method=op.method
+        )
+        rng["jax_state"] = rng_key
+        return (rng, sample)
+
+    return sample_fn
+
+
 @jax_sample_fn.register(ptr.BernoulliRV)
 def jax_sample_fn_bernoulli(op, node):
     """JAX implementation of `BernoulliRV`."""

pytensor/tensor/random/basic.py

@@ -1,13 +1,16 @@
 import abc
 import warnings
+from typing import Literal
 
 import numpy as np
 import scipy.stats as stats
 from numpy import broadcast_shapes as np_broadcast_shapes
 from numpy import einsum as np_einsum
+from numpy import sqrt as np_sqrt
 from numpy.linalg import cholesky as np_cholesky
+from numpy.linalg import eigh as np_eigh
+from numpy.linalg import svd as np_svd
 
-import pytensor
 from pytensor.tensor import get_vector_length, specify_shape
 from pytensor.tensor.basic import as_tensor_variable
 from pytensor.tensor.math import sqrt
@@ -852,8 +855,17 @@ class MvNormalRV(RandomVariable):
     signature = "(n),(n,n)->(n)"
     dtype = "floatX"
     _print_name = ("MultivariateNormal", "\\operatorname{MultivariateNormal}")
+    __props__ = ("name", "signature", "dtype", "inplace", "method")
 
-    def __call__(self, mean=None, cov=None, size=None, **kwargs):
+    def __init__(self, *args, method: Literal["cholesky", "svd", "eigh"], **kwargs):
+        super().__init__(*args, **kwargs)
+        if method not in ("cholesky", "svd", "eigh"):
+            raise ValueError(
+                f"Unknown method {method}. The method must be one of 'cholesky', 'svd', or 'eigh'."
+            )
+        self.method = method
+
+    def __call__(self, mean, cov, size=None, **kwargs):
         r""" "Draw samples from a multivariate normal distribution.
 
         Signature
@@ -876,33 +888,34 @@ def __call__(self, mean=None, cov=None, size=None, **kwargs):
             is specified, a single `N`-dimensional sample is returned.
 
         """
-        dtype = pytensor.config.floatX if self.dtype == "floatX" else self.dtype
-
-        if mean is None:
-            mean = np.array([0.0], dtype=dtype)
-        if cov is None:
-            cov = np.array([[1.0]], dtype=dtype)
         return super().__call__(mean, cov, size=size, **kwargs)
 
-    @classmethod
-    def rng_fn(cls, rng, mean, cov, size):
+    def rng_fn(self, rng, mean, cov, size):
         if size is None:
             size = np_broadcast_shapes(mean.shape[:-1], cov.shape[:-2])
 
-        chol = np_cholesky(cov)
+        if self.method == "cholesky":
+            A = np_cholesky(cov)
+        elif self.method == "svd":
+            A, s, _ = np_svd(cov)
+            A *= np_sqrt(s, out=s)[..., None, :]
+        else:
+            w, A = np_eigh(cov)
+            A *= np_sqrt(w, out=w)[..., None, :]
+
         out = rng.normal(size=(*size, mean.shape[-1]))
         np_einsum(
             "...ij,...j->...i",  # numpy doesn't have a batch matrix-vector product
-            chol,
+            A,
             out,
-            out=out,
             optimize=False,  # Nothing to optimize with two operands, skip costly setup
+            out=out,
         )
         out += mean
         return out
 
 
-multivariate_normal = MvNormalRV()
+multivariate_normal = MvNormalRV(method="cholesky")
 
 
 class DirichletRV(RandomVariable):

tests/tensor/random/test_basic.py

@@ -19,6 +19,7 @@
 from pytensor.tensor import ones, stack
 from pytensor.tensor.random.basic import (
     ChoiceWithoutReplacement,
+    MvNormalRV,
     PermutationRV,
     _gamma,
     bernoulli,
@@ -668,6 +669,36 @@ def test_mvnormal_ShapeFeature():
     assert s4.get_test_value() == 3
 
 
+@pytest.mark.parametrize("psd", (True, False))
+@pytest.mark.parametrize("method", ("cholesky", "svd", "eigh"))
+def test_mvnormal_cov_decomposition_method(method, psd):
+    mean = 2 ** np.arange(3)
+    if psd:
+        cov = [
+            [1, 0.5, -1],
+            [0.5, 2, 0],
+            [-1, 0, 3],
+        ]
+    else:
+        cov = [
+            [1, 0.5, 0],
+            [0.5, 2, 0],
+            [0, 0, 0],
+        ]
+    rng = shared(np.random.default_rng(675))
+    draws = MvNormalRV(method=method)(mean, cov, rng=rng, size=(10_000,))
+    assert draws.owner.op.method == method
+
+    if not psd and method == "cholesky":
+        with pytest.raises(np.linalg.LinAlgError):
+            draws.eval()
+        return
+
+    draws_eval = draws.eval()
+    np.testing.assert_allclose(np.mean(draws_eval, axis=0), mean, rtol=0.01)
+    np.testing.assert_allclose(np.cov(draws_eval, rowvar=False), cov, atol=0.1)
+
+
 @pytest.mark.parametrize(
     "alphas, size",
     [