Repository navigation
find_MAP: jitter has multiple issues #687
Description
find_MAP:
a) jitter_rvs=None appears to not jitter variables (while docstring says None -> All)
b) even with explicit jitter_rvs=model.free_RVs does not work with default freeze_model=True
In factor-model-like cases, this can leave optimization stuck at a bad zero stationary point.
Minimal example:
import pymc as pm
import pymc_extras as pmx
def run(label, **kwargs):
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)
if kwargs.get("jitter_rvs") == "free_RVs":
kwargs["jitter_rvs"] = list(model.free_RVs)
idata = pmx.find_MAP(
model=model,
progressbar=False,
maxiter=1000,
maxfun=1000,
random_seed=11,
**kwargs,
)
w_hat = float(idata.posterior["w"].values[0, 0])
z_hat = float(idata.posterior["z"].values[0, 0])
print(label, w_hat, z_hat, w_hat * z_hat)
run("default")
run("jitter_rvs=None, freeze_model=False", freeze_model=False)
run("jitter_rvs=free_RVs, freeze_model=False", jitter_rvs="free_RVs", freeze_model=False)Observed:
default 0.0 0.0 0.0
jitter_rvs=None, freeze_model=False 0.0 0.0 0.0
jitter_rvs=free_RVs, freeze_model=False 0.995 0.995 0.99
Expected: jitter_rvs=None should jitter all free variables as documented (as pm.sample usually does, assuming this is meant as a drop-in replacement)
Activity
ricardoV94 commentedon Oct 8, 2026
This issues should be closed by redirecting pymc-extras find_MAP to the one merged in PyTensor with a FutureWarning that it will be removed afterwards