refactor init_jitter inputs

Author: Goose

pymc/sampling/jax.py

@@ -213,7 +213,7 @@ def _get_batched_jittered_initial_points(
     chains: int,
     initvals: StartDict | Sequence[StartDict | None] | None,
     random_seed: RandomSeed,
-    logp_fn: Callable[[Sequence[np.ndarray]], np.ndarray],
+    logp_fn: Callable[[Sequence[np.ndarray]], np.ndarray] | None = None,
     jitter: bool = True,
     jitter_max_retries: int = 10,
 ) -> np.ndarray | list[np.ndarray]:
@@ -230,14 +230,18 @@ def _get_batched_jittered_initial_points(
         list with one item per variable and number of chains as batch dimension.
         Each item has shape `(chains, *var.shape)`
     """
+    if logp_fn is None:
+        eval_logp_initial_point = None
+
+    else:
 
-    def eval_logp_initial_point(point: dict[str, np.ndarray]) -> np.ndarray:
-        """Wrap logp_fn to conform to _init_jitter logic.
+        def eval_logp_initial_point(point: dict[str, np.ndarray]) -> np.ndarray:
+            """Wrap logp_fn to conform to _init_jitter logic.
 
-        Wraps jaxified logp function to accept a dict of
-        {model_variable: np.array} key:value pairs.
-        """
-        return logp_fn(point.values())
+            Wraps jaxified logp function to accept a dict of
+            {model_variable: np.array} key:value pairs.
+            """
+            return logp_fn(point.values())
 
     initial_points = _init_jitter(
         model,

pymc/sampling/mcmc.py

@@ -1338,7 +1338,6 @@ def _init_jitter(
     seeds: Sequence[int] | np.ndarray,
     jitter: bool,
     jitter_max_retries: int,
-    logp_dlogp_func=None,
     logp_fn: Callable[[PointType], np.ndarray] | None = None,
 ) -> list[PointType]:
     """Apply a uniform jitter in [-1, 1] to the test value as starting point in each chain.
@@ -1354,8 +1353,9 @@ def _init_jitter(
         Whether to apply jitter or not.
     jitter_max_retries : int
         Maximum number of repeated attempts at initializing values (per chain).
-    logp_fn: Callable[[dict[str, np.ndarray]], np.ndarray]
+    logp_fn: Callable[[dict[str, np.ndarray]], np.ndarray] | None
         Jaxified logp function that takes the output of the initial point functions as input.
+        If None, will use the results of model.compile_logp().
 
     Returns
     -------
@@ -1372,19 +1372,10 @@ def _init_jitter(
     if not jitter:
         return [ipfn(seed) for ipfn, seed in zip(ipfns, seeds)]
 
-    model_logp_fn: Callable[[PointType], np.ndarray]
-    if logp_dlogp_func is None:
-        if logp_fn is None:
-            # pymc NUTS path
-            model_logp_fn = model.compile_logp()
-        else:
-            # Jax path
-            model_logp_fn = logp_fn
+    if logp_fn is None:
+        model_logp_fn = model.compile_logp()
     else:
-
-        def model_logp_fn(ip: PointType) -> np.ndarray:
-            q, _ = DictToArrayBijection.map(ip)
-            return logp_dlogp_func([q], extra_vars={})[0]
+        model_logp_fn = logp_fn
 
     initial_points = []
     for ipfn, seed in zip(ipfns, seeds):
@@ -1509,13 +1500,18 @@ def init_nuts(
 
     logp_dlogp_func = model.logp_dlogp_function(ravel_inputs=True, **compile_kwargs)
     logp_dlogp_func.trust_input = True
+
+    def model_logp_fn(ip: PointType) -> np.ndarray:
+        q, _ = DictToArrayBijection.map(ip)
+        return logp_dlogp_func([q], extra_vars={})[0]
+
     initial_points = _init_jitter(
         model,
         initvals,
         seeds=random_seed_list,
         jitter="jitter" in init,
         jitter_max_retries=jitter_max_retries,
-        logp_dlogp_func=logp_dlogp_func,
+        logp_fn=model_logp_fn,
     )
 
     apoints = [DictToArrayBijection.map(point) for point in initial_points]