Add jitter_scale parameter for initial point generation

pymc/initial_point.py

@@ -71,8 +71,7 @@ def make_initial_point_fns_per_chain(
 ) -> list[Callable]:
     """Create an initial point function for each chain, as defined by initvals.
 
-    If a single initval dictionary is passed, the function is replicated for each
-    chain, otherwise a unique function is compiled for each entry in the dictionary.
+    If a single initval dictionary is passed, the function is replicated for each chain, otherwise a unique function is compiled for each entry in the dictionary.
 
     Parameters
     ----------
@@ -82,6 +81,8 @@ def make_initial_point_fns_per_chain(
     jitter_rvs : set, optional
         Random variable tensors for which U(-1, 1) jitter shall be applied.
         (To the transformed space if applicable.)
+    jitter_scale : float, optional
+        The scale of the jitter in the jitter_rvs set. Defaults to 1.0.
 
     Raises
     ------
@@ -134,8 +135,9 @@ def make_initial_point_fn(
     Parameters
     ----------
     jitter_rvs : set
-        The set (or list or tuple) of random variables for which a U(-1, +1) jitter should be
-        added to the initial value. Only available for variables that have a transform or real-valued support.
+        The set (or list or tuple) of random variables for which a U(-1, +1) jitter should be added to the initial value. Only available for variables that have a transform or real-valued support.
+    jitter_scale : float, optional
+        The scale of the jitter in the jitter_rvs set. Defaults to 1.0.
     default_strategy : str
         Which of { "support_point", "prior" } to prefer if the initval setting for an RV is None.
     overrides : dict
@@ -209,8 +211,10 @@ def make_initial_point_expression(
         Mapping of free random variable tensors to initial value strategies.
         For example the `Model.initial_values` dictionary.
     jitter_rvs : set
-        The set (or list or tuple) of random variables for which a U(-1, +1) jitter should be
+        The set (or list or tuple) of random variables for which a U(-1, 1) jitter should be
         added to the initial value. Only available for variables that have a transform or real-valued support.
+    jitter_scale : float, optional
+        The scale of the jitter in the jitter_rvs set. Defaults to 1.0.
     default_strategy : str
         Which of { "support_point", "prior" } to prefer if the initval strategy setting for an RV is None.
     return_transformed : bool

pymc/sampling/mcmc.py

@@ -1438,6 +1438,8 @@ def _init_jitter(
     ----------
     jitter: bool
         Whether to apply jitter or not.
+    jitter_scale : float, optional
+        The scale of the jitter in set(model.free_RVs). Defaults to 1.0.
     jitter_max_retries : int
         Maximum number of repeated attempts at initializing values (per chain).
 

tests/test_initial_point.py

@@ -156,29 +156,20 @@ def test_jitter_scale(self):
         with pm.Model() as pmodel:
             A = pm.HalfFlat("A", initval="support_point")
 
-        jitter_scale_tests = np.array([1.0, 2.0, 5.0])
-        fns = []
-        for jitter_scale in jitter_scale_tests:
-            fns.append(
-                make_initial_point_fn(
-                    model=pmodel,
-                    jitter_rvs=set(pmodel.free_RVs),
-                    jitter_scale=jitter_scale,
-                    return_transformed=True,
-                )
-            )
-
-        n_draws = 1000
-        jitter_samples = np.empty((n_draws, len(fns)))
-        for j, fn in enumerate(fns):
-            # start and end to ensure random samples, otherwise jitter_samples across different jitter_scale will be an exact scale of each other
-            start = j * n_draws
-            end = start + n_draws
-            jitter_samples[:, j] = np.asarray([fn(i)["A_log__"] for i in range(start, end)])
+        fn_default = make_initial_point_fn(
+            model=pmodel,
+            jitter_rvs=set(pmodel.free_RVs),
+            return_transformed=True,
+        )
 
-        init_standardised = np.mean((jitter_samples / jitter_scale_tests), axis=0)
+        fn_large = make_initial_point_fn(
+            model=pmodel,
+            jitter_rvs=set(pmodel.free_RVs),
+            jitter_scale=1000.0,
+            return_transformed=True,
+        )
 
-        assert np.all((-0.05 < init_standardised) & (init_standardised < 0.05))
+        assert fn_large(0)["A_log__"] > fn_default(0)["A_log__"]
 
     def test_respects_overrides(self):
         with pm.Model() as pmodel: