Fix jitter_rvs handling in find_MAP

pymc_extras/inference/laplace_approx/find_map.py

@@ -56,7 +56,12 @@ def get_nearest_psd(A: np.ndarray) -> np.ndarray:
 
 
 def _make_initial_point(model, initvals=None, random_seed=None, jitter_rvs=None):
-    jitter_rvs = [] if jitter_rvs is None else jitter_rvs
+    if jitter_rvs is None:
+        jitter_rvs = model.free_RVs
+    else:
+        # The model can be a frozen copy of the one the variables were taken from, so match them by name
+        jitter_names = {rv.name for rv in jitter_rvs}
+        jitter_rvs = [rv for rv in model.free_RVs if rv.name in jitter_names]
 
     ipfn = make_initial_point_fn(
         model=model,

tests/inference/laplace_approx/test_find_map.py

@@ -353,3 +353,36 @@ def test_map_nonscalar_rv_without_dims():
         "x[1,1]",
         "x[1,2]",
     ]
+
+
+@pytest.mark.parametrize("freeze_model", [True, False], ids=["frozen", "not_frozen"])
+@pytest.mark.parametrize("jitter_rvs", [None, "all", "w"], ids=str)
+def test_find_MAP_jitter_rvs(jitter_rvs, freeze_model):
+    # Without jitter, the optimizer starts at w = z = 0, a stationary point it never leaves
+    with pm.Model() as model:
+        w = pm.Normal("w")
+        z = pm.Normal("z")
+        pm.Normal("y", mu=w * z, sigma=0.1, observed=1.0)
+
+    jitter_rvs = {None: None, "all": [w, z], "w": [w]}[jitter_rvs]
+    idata = find_MAP(
+        model=model,
+        jitter_rvs=jitter_rvs,
+        freeze_model=freeze_model,
+        random_seed=11,
+        progressbar=False,
+    )
+    w_hat = idata.posterior["w"].item()
+    z_hat = idata.posterior["z"].item()
+    assert w_hat * z_hat == pytest.approx(0.99, abs=1e-3)
+
+
+def test_find_MAP_no_jitter():
+    with pm.Model() as model:
+        w = pm.Normal("w")
+        z = pm.Normal("z")
+        pm.Normal("y", mu=w * z, sigma=0.1, observed=1.0)
+
+    idata = find_MAP(model=model, jitter_rvs=[], random_seed=11, progressbar=False)
+    assert idata.posterior["w"].item() == 0.0
+    assert idata.posterior["z"].item() == 0.0