Added metropolis hastings sampling for jax

Author: FlorianPfaff

.pylintrc

@@ -63,7 +63,7 @@ ignore-patterns=^\.#
 # (useful for modules/projects where namespaces are manipulated during runtime
 # and thus existing member attributes cannot be deduced by static analysis). It
 # supports qualified module names, as well as Unix pattern matching.
-ignored-modules=pyrecest.backend
+ignored-modules=pyrecest.backend, jax
 
 # Python code to execute, usually for sys.path manipulation such as
 # pygtk.require().

pyrecest/distributions/abstract_manifold_specific_distribution.py

@@ -1,10 +1,11 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from typing import Union
+import inspect
 
 import pyrecest.backend
 
-# pylint: disable=no-name-in-module,no-member
+# pylint: disable=no-name-in-module,no-member,redefined-builtin
 from pyrecest.backend import empty, int32, int64, log, random, squeeze
 
 
@@ -64,13 +65,14 @@ def set_mode(self, _):
         """
         raise NotImplementedError("set_mode is not implemented for this distribution")
 
+    # Need to use Union instead of | to support torch.dtype
     # Need to use Union instead of | to support torch.dtype
     def sample(self, n: Union[int, int32, int64]):
         """Obtain n samples from the distribution."""
         return self.sample_metropolis_hastings(n)
 
     # jscpd:ignore-start
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def sample_metropolis_hastings(
         self,
         n: Union[int, int32, int64],
@@ -81,30 +83,48 @@ def sample_metropolis_hastings(
     ):
         # jscpd:ignore-end
         """Metropolis Hastings sampling algorithm."""
-        assert (
-            pyrecest.backend.__backend_name__ != "jax"
-        ), "Not supported on this backend"
+        if pyrecest.backend.__backend_name__ == "jax":
+            # Get a key from your global JAX random state *outside* of lax.scan
+            import jax as _jax
+
+            key = random.jax_global_random_state()
+            key, key_for_mh = _jax.random.split(key)
+            # Optionally update global state for future calls
+            random.jax_global_random_state(key)
+
+            if proposal is None or start_point is None:
+                raise NotImplementedError(
+                    "Default proposals and starting points should be set in inheriting classes."
+                )
+            _assert_proposal_supports_key(proposal)
+
+            samples, _ = sample_metropolis_hastings_jax(
+                key=key_for_mh,
+                log_pdf=self.ln_pdf,
+                proposal=proposal,  # must be (key, x) -> x_prop for JAX
+                start_point=start_point,
+                n=int(n),
+                burn_in=int(burn_in),
+                skipping=int(skipping),
+            )
+            # You could optionally stash `key_out` somewhere if you want chain continuation.
+            return squeeze(samples)
+
+        # Non-JAX backends → your old NumPy/Torch code
         if proposal is None or start_point is None:
             raise NotImplementedError(
                 "Default proposals and starting points should be set in inheriting classes."
             )
 
         total_samples = burn_in + n * skipping
-        s = empty(
-            (
-                total_samples,
-                self.input_dim,
-            ),
-        )
+        s = empty((total_samples, self.input_dim))
         x = start_point
         i = 0
         pdfx = self.pdf(x)
 
         while i < total_samples:
             x_new = proposal(x)
-            assert (
-                x_new.shape == x.shape
-            ), "Proposal must return a vector of same shape as input"
+            assert x_new.shape == x.shape, "Proposal must return a vector of same shape as input"
             pdfx_new = self.pdf(x_new)
             a = pdfx_new / pdfx
             if a.item() > 1 or a.item() > random.rand(1):
@@ -115,3 +135,111 @@ def sample_metropolis_hastings(
 
         relevant_samples = s[burn_in::skipping, :]
         return squeeze(relevant_samples)
+
+# pylint: disable=too-many-positional-arguments,too-many-locals,too-many-arguments
+def sample_metropolis_hastings_jax(
+    key,
+    log_pdf,      # function: x -> log p(x)
+    proposal,     # function: (key, x) -> x_prop
+    start_point,
+    n: int,
+    burn_in: int = 10,
+    skipping: int = 5,
+):
+    """
+    Metropolis-Hastings sampler in JAX.
+
+    key:        jax.random.PRNGKey
+    log_pdf:    callable x -> log p(x)
+    proposal:   callable (key, x) -> x_proposed
+    start_point: initial state (array)
+    n:          number of samples to return (after burn-in and thinning)
+    """
+    import jax.numpy as _jnp
+    from jax import _random, _lax
+
+
+    start_point = _jnp.asarray(start_point)
+    total_steps = burn_in + n * skipping
+
+    def one_step(carry, _):
+        key, x, log_px = carry
+        key, key_prop, key_u = _random.split(key, 3)
+
+        # Propose new state
+        x_prop = proposal(key_prop, x)
+        log_px_prop = log_pdf(x_prop)
+
+        # log_alpha = log p(x_prop) - log p(x)
+        log_alpha = log_px_prop - log_px
+
+        # Draw u ~ Uniform(0, 1)
+        u = _random.uniform(key_u, shape=())
+        log_u = _jnp.log(u)
+
+        # Accept if log u < min(0, log_alpha)
+        # (equivalent to u < exp(min(0, log_alpha)))
+        log_alpha_capped = _jnp.minimum(0.0, log_alpha)
+        accept = log_u < log_alpha_capped  # scalar bool
+
+        # Branch without Python if
+        x_new = _jnp.where(accept, x_prop, x)
+        log_px_new = _jnp.where(accept, log_px_prop, log_px)
+
+        return (key, x_new, log_px_new), x_new
+
+    init_carry = (key, start_point, log_pdf(start_point))
+    (key_out, _, _), chain = _lax.scan(
+        one_step,
+        init_carry,
+        xs=None,
+        length=total_steps,
+    )
+
+    samples = chain[burn_in::skipping]
+    return samples, key_out
+
+
+def _assert_proposal_supports_key(proposal: Callable):
+    """
+    Check that `proposal` can be called as proposal(key, x).
+
+    Raises a TypeError with a helpful message if this is not the case.
+    """
+    # Unwrap jitted / partial / decorated functions if possible
+    func = proposal
+    while hasattr(func, "__wrapped__"):
+        func = func.__wrapped__
+
+    try:
+        sig = inspect.signature(func)
+    except (TypeError, ValueError):
+        # Can't introspect (e.g. builtins); fall back to a generic error
+        raise TypeError(
+            "For the JAX backend, `proposal` must accept (key, x) as arguments, "
+            "but its signature could not be inspected."
+        ) from None
+
+    params = list(sig.parameters.values())
+
+    # Count positional(-or-keyword) parameters
+    num_positional = sum(
+        p.kind in (inspect.Parameter.POSITIONAL_ONLY,
+                   inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        for p in params
+    )
+    has_var_positional = any(
+        p.kind == inspect.Parameter.VAR_POSITIONAL
+        for p in params
+    )
+
+    if has_var_positional or num_positional >= 2:
+        # Looks compatible with (key, x)
+        return
+
+    raise TypeError(
+        "For the JAX backend, `proposal` must accept `(key, x)` as arguments.\n"
+        f"Got signature: {sig}\n"
+        "Hint: change your proposal from `def proposal(x): ...` to\n"
+        "`def proposal(key, x): ...` and use `jax.random` with the passed key."
+    )

pyrecest/tests/distributions/test_abstract_mixture.py

@@ -25,10 +25,6 @@ def _test_sample(self, mix, n):
             self.assertEqual(s.shape, (n, mix.input_dim))
         return s
 
-    @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
-        reason="Not supported on this backend",
-    )
     def test_sample_metropolis_hastings_basics_only_t2(self):
         vmf = ToroidalWrappedNormalDistribution(array([1.0, 0.0]), eye(2))
         mix = HypertoroidalMixture(
@@ -37,28 +33,28 @@ def test_sample_metropolis_hastings_basics_only_t2(self):
         self._test_sample(mix, 10)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+        pyrecest.backend.__backend_name__ in ("pytorch",),
         reason="Not supported on this backend",
     )
     def test_sample_metropolis_hastings_basics_only_s2(self):
         vmf1 = VonMisesFisherDistribution(
             array([1.0, 0.0, 0.0]), 2.0
-        )  # Needs to be float for scipy
+        )
         vmf2 = VonMisesFisherDistribution(
             array([0.0, 1.0, 0.0]), 2.0
-        )  # Needs to be float for scipy
+        )
         mix = HypersphericalMixture([vmf1, vmf2], array([0.5, 0.5]))
         s = self._test_sample(mix, 10)
         self.assertTrue(allclose(linalg.norm(s, axis=1), ones(10), rtol=1e-10))
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+        pyrecest.backend.__backend_name__ in ("pytorch",),
         reason="Not supported on this backend",
     )
     def test_sample_metropolis_hastings_basics_only_h2(self):
         vmf = VonMisesFisherDistribution(
             array([1.0, 0.0, 0.0]), 2.0
-        )  # Needs to be float for scipy
+        )
         mix = CustomHyperhemisphericalDistribution(
             lambda x: vmf.pdf(x) + vmf.pdf(-x), 2
         )