Open
Description
I'm using Pathfinder to generate inits for a subsequent run of NUTS, and for each of my N
NUTS chains, I run a separate multi-path pathfinder:
n_chains = int(os.cpu_count()/2)
init_vals = []
seed = 0
while len(init_vals) < n_chains:
try:
seed += 1
with inference_model:
pathfinder_trace = pmx.fit(
method = "pathfinder"
, random_seed = seed
, num_paths = 5 # more than 5 causes a crash!
, num_draws_per_path = 1000
, importance_sampling = 'psis'
, num_draws = 1
)
pathfinder_trace.posterior = pathfinder_trace.posterior.squeeze("chain", drop=True)
point = {
var: pathfinder_trace.posterior[var].isel(draw=0).values
for var in pathfinder_trace.posterior.data_vars
}
init_vals.append(point)
except Exception as e:
print(f"Error occurred: {e}")
print("Retrying with a different seed...")
# now we can use the pathfinder trace to initialize the chains
with inference_model:
posterior_trace = pm.sample(
chains = n_chains
, cores = n_chains
, initvals = init_vals
)
However, with the current API, there is an unnecessary re-compilation of the model for every run of pathfinder, so it would be useful to have separate compile()
and run()
methods.
Metadata
Metadata
Assignees
Labels
No labels