Skip to content

FR: Separate compile/run methods for pathfinder #488

Open
@mike-lawrence

Description

@mike-lawrence

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

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions