Add JAX implementation fol MatrixIsPositiveDefinite Op

pymc/distributions/multivariate.py

@@ -849,9 +849,9 @@
 def posdef(AA):
     try:
         linalg.cholesky(AA)
-        return 1
+        return True
     except linalg.LinAlgError:
-        return 0
+        return False
 
 
 class PosDefMatrix(Op):
@@ -868,15 +868,15 @@
     def make_node(self, x):
         x = pt.as_tensor_variable(x)
         assert x.ndim == 2
-        o = TensorType(dtype="int8", shape=[])()
+        o = TensorType(dtype="bool", shape=[])()
         return Apply(self, [x], [o])
 
     # Python implementation:
     def perform(self, node, inputs, outputs):
         (x,) = inputs
         (z,) = outputs
         try:
-            z[0] = np.array(posdef(x), dtype="int8")
+            z[0] = np.array(posdef(x), dtype="bool")
         except Exception:
             pm._log.exception("Failed to check if %s positive definite", x)
             raise

pymc/sampling/jax.py

@@ -21,6 +21,7 @@
 
 import arviz as az
 import jax
+import jax.numpy as jnp
 import numpy as np
 import pytensor.tensor as pt
 
@@ -34,10 +35,10 @@
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.type import RandomType
-from pytensor.tensor.shape import SpecifyShape
 
 from pymc import Model, modelcontext
 from pymc.backends.arviz import find_constants, find_observations
+from pymc.distributions.multivariate import PosDefMatrix
 from pymc.initial_point import StartDict
 from pymc.logprob.utils import CheckParameterValue
 from pymc.sampling.mcmc import _init_jitter
@@ -62,7 +63,6 @@
 
 @jax_funcify.register(Assert)
 @jax_funcify.register(CheckParameterValue)
-@jax_funcify.register(SpecifyShape)
 def jax_funcify_Assert(op, **kwargs):
     # Jax does not allow assert whose values aren't known during JIT compilation
     # within it's JIT-ed code. Hence we need to make a simple pass through
@@ -74,6 +74,15 @@ def assert_fn(value, *inps):
     return assert_fn
 
 
+@jax_funcify.register(PosDefMatrix)
+def jax_funcify_PosDefMatrix(op, **kwargs):
+    def posdefmatrix_fn(value, *inps):
+        no_pos_def = jnp.any(jnp.isnan(jnp.linalg.cholesky(value)))
+        return jnp.invert(no_pos_def)
+
+    return posdefmatrix_fn
+
+
 def _replace_shared_variables(graph: List[TensorVariable]) -> List[TensorVariable]:
     """Replace shared variables in graph by their constant values
 

tests/distributions/test_multivariate.py

@@ -1907,26 +1907,22 @@ def check_draws(self):
         def ref_rand(mu, rowcov, colcov):
             return st.matrix_normal.rvs(mean=mu, rowcov=rowcov, colcov=colcov)
 
-        with pm.Model():
-            matrixnormal = pm.MatrixNormal(
-                "matnormal",
-                mu=np.random.random((3, 3)),
-                rowcov=np.eye(3),
-                colcov=np.eye(3),
-            )
-            check = pm.sample_prior_predictive(n_fails, return_inferencedata=False, random_seed=1)
-
-        ref_smp = ref_rand(mu=np.random.random((3, 3)), rowcov=np.eye(3), colcov=np.eye(3))
+        matrixnormal = pm.MatrixNormal.dist(
+            mu=np.random.random((3, 3)),
+            rowcov=np.eye(3),
+            colcov=np.eye(3),
+        )
 
         p, f = delta, n_fails
         while p <= delta and f > 0:
-            matrixnormal_smp = check["matnormal"]
+            matrixnormal_smp = pm.draw(matrixnormal)
+            ref_smp = ref_rand(mu=np.random.random((3, 3)), rowcov=np.eye(3), colcov=np.eye(3))
 
             p = np.min(
                 [
                     st.ks_2samp(
-                        np.atleast_1d(matrixnormal_smp).flatten(),
-                        np.atleast_1d(ref_smp).flatten(),
+                        matrixnormal_smp.flatten(),
+                        ref_smp.flatten(),
                     )
                 ]
             )
@@ -2134,10 +2130,10 @@ def test_car_rng_fn(sparse):
 @pytest.mark.parametrize(
     "matrix, result",
     [
-        ([[1.0, 0], [0, 1]], 1),
-        ([[1.0, 2], [2, 1]], 0),
-        ([[1.0, 1], [1, 1]], 0),
-        ([[1, 0.99, 1], [0.99, 1, 0.999], [1, 0.999, 1]], 0),
+        ([[1.0, 0], [0, 1]], True),
+        ([[1.0, 2], [2, 1]], False),
+        ([[1.0, 1], [1, 1]], False),
+        ([[1, 0.99, 1], [0.99, 1, 0.999], [1, 0.999, 1]], False),
     ],
 )
 def test_posdef_symmetric(matrix, result):

tests/gp/test_hsgp_approx.py

@@ -166,7 +166,7 @@ def test_parametrization_drop_first(self, model, cov_func, X1, drop_first):
                 assert n_coeffs == n_basis, "one was dropped when it shouldn't have been"
 
     @pytest.mark.parametrize("parameterization", ["centered", "noncentered"])
-    def test_prior(self, model, cov_func, X1, parameterization):
+    def test_prior(self, model, cov_func, X1, parameterization, rng):
         """Compare HSGP prior to unapproximated GP prior, pm.gp.Latent.  Draw samples from the
         prior and compare them using MMD two sample test.  Tests both centered and non-centered
         parameterizations.
@@ -178,7 +178,7 @@ def test_prior(self, model, cov_func, X1, parameterization):
             gp = pm.gp.Latent(cov_func=cov_func)
             f2 = gp.prior("f2", X=X1)
 
-            idata = pm.sample_prior_predictive(samples=1000)
+            idata = pm.sample_prior_predictive(samples=1000, random_seed=rng)
 
         samples1 = az.extract(idata.prior["f1"])["f1"].values.T
         samples2 = az.extract(idata.prior["f2"])["f2"].values.T

tests/sampling/test_jax.py

@@ -18,6 +18,7 @@
 
 import arviz as az
 import jax
+import jax.numpy as jnp
 import numpy as np
 import pytensor
 import pytensor.tensor as pt
@@ -29,6 +30,7 @@
 
 import pymc as pm
 
+from pymc.distributions.multivariate import PosDefMatrix
 from pymc.sampling.jax import (
     _get_batched_jittered_initial_points,
     _get_log_likelihood,
@@ -49,6 +51,25 @@ def test_old_import_route():
     assert set(new_sj.__all__) <= set(dir(old_sj))
 
 
+def test_jax_PosDefMatrix():
+    x = pt.tensor(name="x", shape=(2, 2), dtype="float32")
+    matrix_pos_def = PosDefMatrix()
+    x_is_pos_def = matrix_pos_def(x)
+    f = pytensor.function(inputs=[x], outputs=[x_is_pos_def], mode="JAX")
+
+    test_cases = [
+        (jnp.eye(2), True),
+        (jnp.zeros(shape=(2, 2)), False),
+        (jnp.array([[1, -1.5], [0, 1.2]], dtype="float32"), True),
+        (-1 * jnp.array([[1, -1.5], [0, 1.2]], dtype="float32"), False),
+        (jnp.array([[1, -1.5], [0, -1.2]], dtype="float32"), False),
+    ]
+
+    for input, expected in test_cases:
+        actual = f(input)[0]
+        assert jnp.array_equal(a1=actual, a2=expected)
+
+
 @pytest.mark.parametrize(
     "sampler",
     [