Add JAX implementation of PosDefMatrix Op

Author: juanitorduz

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
 
@@ -37,6 +38,7 @@
 
 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
@@ -72,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/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",
     [