Sharding
Multi-device strategies for BlackJAX nested sampling.
Functions:
| Name | Description |
|---|---|
build_sharded_from_mcmc_kernel |
Assemble BlackJAX's adaptive NS kernel with sharded strategies. |
delete_fn_sharded |
Select the globally lowest-likelihood live particles. |
make_live_mesh |
Build a single-host device mesh for the live-particle axis. |
place_key |
Replicate a sampler PRNG key over the mesh. |
place_state |
Shard live particles and replicate adaptive/integrator state. |
update_with_mcmc_take_last_sharded |
Run independent constrained replacement chains across the mesh. |
build_sharded_from_mcmc_kernel(constrained_step_fn: Callable, n_inner_steps: int, update_inner_kernel_params_fn: Callable, n_delete: int, mesh: Mesh) -> Callable
¤
Assemble BlackJAX's adaptive NS kernel with sharded strategies.
delete_fn_sharded(mesh: Mesh, n_delete: int) -> Callable
¤
Select the globally lowest-likelihood live particles.
make_live_mesh(n_devices: int, n_live: int, n_delete: int) -> Optional[Mesh]
¤
Build a single-host device mesh for the live-particle axis.
place_key(rng_key, mesh: Mesh)
¤
Replicate a sampler PRNG key over the mesh.
place_state(state, mesh: Mesh)
¤
Shard live particles and replicate adaptive/integrator state.
update_with_mcmc_take_last_sharded(constrained_step_fn: Callable, n_inner_steps: int, n_delete: int, mesh: Mesh) -> Callable
¤
Run independent constrained replacement chains across the mesh.