diff --git a/docs/guides/cli.md b/docs/guides/cli.md index 4d50b993..2916c2eb 100644 --- a/docs/guides/cli.md +++ b/docs/guides/cli.md @@ -354,6 +354,7 @@ n_tempered_steps = 5 | `n_delete_frac` | `0.5` | Fraction of live points replaced per iteration | | `num_inner_steps_per_dim` | `10` | Slice steps per dimension for the nested kernel | | `termination_dlogz` | `0.1` | Stop when the remaining log-evidence contribution falls below this | +| `n_devices` | `1` | Number of local devices used to shard live points; `n_live` and `n_delete` must be divisible by it | | `checkpoint_dir` | `{output.dir}/` | Directory for `checkpoint.pkl`; set by the CLI automatically | | `checkpoint_interval` | `600.0` | Seconds between checkpoint writes; `0` disables checkpointing | @@ -369,6 +370,7 @@ n_tempered_steps = 5 | `max_steps` | `10` | Maximum stepping-out expansions per slice | | `max_shrinkage` | `100` | Maximum shrinkage evaluations per slice | | `termination_dlogz` | `exp(-3)` (~`0.0498`) | Stop when the remaining evidence contribution falls below this | +| `n_devices` | `1` | Number of local devices used to shard live points; `n_live` and `n_delete` must be divisible by it | | `checkpoint_dir` | `{output.dir}/` | Directory for `checkpoint.pkl`; set by the CLI automatically | | `checkpoint_interval` | `600.0` | Seconds between checkpoint writes; `0` disables checkpointing | diff --git a/docs/guides/samplers.md b/docs/guides/samplers.md index 206fda07..3ea38a48 100644 --- a/docs/guides/samplers.md +++ b/docs/guides/samplers.md @@ -149,6 +149,7 @@ jim = Jim( n_delete_frac=0.5, num_inner_steps_per_dim=20, termination_dlogz=0.1, + n_devices=1, ), sample_transforms=sample_transforms, likelihood_transforms=likelihood_transforms, @@ -162,6 +163,9 @@ Key parameters: - `n_live` — number of live points. - `n_delete_frac` — fraction of live points replaced per iteration. - `num_inner_steps_per_dim` — slice-sampler steps per dimension per dead point; increase for strongly correlated posteriors. +- `n_devices` — number of local devices used to shard the live points and + replacement chains. Both `n_live` and `int(n_live * n_delete_frac)` must be + divisible by this value. The default `1` uses the unsharded kernel. **Repository:** [blackjax-devs/blackjax](https://github.com/blackjax-devs/blackjax) @@ -202,6 +206,7 @@ jim = Jim( num_inner_steps_per_dim=1, num_gibbs_sweeps=2, termination_dlogz=math.exp(-3.0), + n_devices=1, ), likelihood_transforms=likelihood_transforms, ) @@ -218,6 +223,9 @@ Key parameters: - `n_live` / `n_delete_frac` — live-set size and replacement fraction, matching NSS. - `num_inner_steps_per_dim` — random-direction slice steps per block dimension. - `num_gibbs_sweeps` — complete block sweeps per particle replacement. +- `n_devices` — number of local devices used to shard the live points and + replacement chains. The waveform cache remains local to each replacement + chain; only live-point state participates in collectives. --- diff --git a/src/jimgw/samplers/blackjax/nss.py b/src/jimgw/samplers/blackjax/nss.py index d77e6363..36b04d3c 100644 --- a/src/jimgw/samplers/blackjax/nss.py +++ b/src/jimgw/samplers/blackjax/nss.py @@ -16,9 +16,22 @@ import blackjax from blackjax.ns.adaptive import AdaptiveNSState, init as _ns_adaptive_init from blackjax.ns.base import NSInfo, init_state_strategy as _init_state_strategy -from blackjax.ns.nss import live_covariance, sample_direction_from_covariance +from blackjax.ns.nss import ( + live_covariance, + sample_direction_from_covariance, + slice_constrained_step, +) from blackjax.ns.utils import finalise +from blackjax.mcmc.slice import build_kernel as build_slice_kernel +from blackjax.mcmc.slice import stepping_out +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P from jimgw.samplers.base import Sampler +from jimgw.samplers.blackjax.sharding import ( + build_sharded_from_mcmc_kernel, + make_live_mesh, + place_key, + place_state, +) from jimgw.samplers.config import BlackJAXNSSConfig from jimgw.samplers.periodic import to_prior_space_proposal @@ -84,9 +97,33 @@ def _checkpoint_tag(self) -> str: def _update_inner_kernel_params_fn(self) -> Callable: return live_covariance - def _build_nested_sampler(self, n_delete: int): + def _build_nested_sampler(self, n_delete: int, mesh: Mesh | None = None): config = self._config num_inner_steps = config.num_inner_steps_per_dim * self.n_dims + if mesh is not None: + init_state_fn = partial( + _init_state_strategy, + logprior_fn=self._log_prior_fn, + loglikelihood_fn=self._log_likelihood_fn, + ) + slice_kernel = build_slice_kernel( + interval=stepping_out, + max_expansions=10, + max_shrinkage=100, + ) + constrained_step = slice_constrained_step( + init_state_fn, slice_kernel, self._proposal + ) + kernel = build_sharded_from_mcmc_kernel( + constrained_step, + n_inner_steps=num_inner_steps, + update_inner_kernel_params_fn=self._update_inner_kernel_params_fn, + n_delete=n_delete, + mesh=mesh, + ) + return blackjax.SamplingAlgorithm( + lambda position, rng_key=None: position, kernel + ) return blackjax.nss( logprior_fn=self._log_prior_fn, loglikelihood_fn=self._log_likelihood_fn, @@ -120,6 +157,7 @@ def _sample( config = self._config n_live = config.n_live n_delete = int(n_live * config.n_delete_frac) + mesh = make_live_mesh(config.n_devices, n_live, n_delete) ckpt_path = ( config.checkpoint_dir / "checkpoint.pkl" if config.checkpoint_dir is not None @@ -137,7 +175,7 @@ def _validated_initial_particles(pos): ) return arr - nested_sampler = self._build_nested_sampler(n_delete) + nested_sampler = self._build_nested_sampler(n_delete, mesh) # Bypass BlackJAX's jax.vmap(init_state_fn) to avoid peak-memory OOM. # A full vmap over all live particles materialises O(n_live) concurrent @@ -151,14 +189,18 @@ def _validated_initial_particles(pos): ) def _batched_nss_init(positions): + if mesh is not None: + positions = jax.device_put(positions, NamedSharding(mesh, P("live"))) + def _batched_fn(pos): return jax.lax.map(_single_init_fn, pos, batch_size=n_delete) - return _ns_adaptive_init( + state = _ns_adaptive_init( positions, init_state_fn=_batched_fn, update_inner_kernel_params_fn=self._update_inner_kernel_params_fn, ) + return place_state(state, mesh) if mesh is not None else state # Resume from checkpoint if one exists. if ( @@ -172,6 +214,9 @@ def _batched_fn(pos): state = _ckpt["state"] dead = _ckpt["dead"] rng_key = _ckpt["rng_key"] + if mesh is not None: + state = place_state(state, mesh) + rng_key = place_key(rng_key, mesh) n_iter = _ckpt["n_iter"] self._prev_elapsed = float(_ckpt["elapsed_time"]) logger.info( @@ -204,6 +249,9 @@ def _batched_fn(pos): dead = [] n_iter = 0 + if mesh is not None: + rng_key = place_key(rng_key, mesh) + def _terminate(state: AdaptiveNSState) -> bool: dlogz = jnp.logaddexp(0, state.integrator.logZ_live - state.integrator.logZ) return bool(jnp.isfinite(dlogz) and dlogz < config.termination_dlogz) @@ -223,9 +271,9 @@ def _terminate(state: AdaptiveNSState) -> bool: ): _last_ckpt_t = config.write_checkpoint( { - "state": state, - "dead": dead, - "rng_key": rng_key, + "state": jax.device_get(state), + "dead": jax.device_get(dead), + "rng_key": jax.device_get(rng_key), "n_iter": n_iter, "elapsed_time": self._prev_elapsed + (time.perf_counter() - _method_t0), @@ -233,7 +281,8 @@ def _terminate(state: AdaptiveNSState) -> bool: self._checkpoint_tag, ) - self._final_state = finalise(state, dead) # type: ignore[arg-type] # AdaptiveNSState structurally satisfies NSState (.particles field) + final_state = finalise(state, dead) # type: ignore[arg-type] # AdaptiveNSState structurally satisfies NSState (.particles field) + self._final_state = jax.device_get(final_state) self._n_iterations = n_iter # Build anesthetic NestedSamples for use in get_samples() and get_diagnostics(). diff --git a/src/jimgw/samplers/blackjax/sharding.py b/src/jimgw/samplers/blackjax/sharding.py new file mode 100644 index 00000000..f815882b --- /dev/null +++ b/src/jimgw/samplers/blackjax/sharding.py @@ -0,0 +1,177 @@ +"""Multi-device strategies for BlackJAX nested sampling.""" + +from __future__ import annotations + +from typing import Callable + +import jax +import jax.numpy as jnp +import numpy as np +from blackjax.ns.adaptive import build_kernel as build_adaptive_kernel +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P + +_LIVE_AXIS = "live" + + +def make_live_mesh(n_devices: int, n_live: int, n_delete: int) -> Mesh | None: + """Build a single-host device mesh for the live-particle axis.""" + if n_devices == 1: + return None + + devices = jax.local_devices() + if n_devices > len(devices): + raise ValueError( + f"n_devices={n_devices} requested, but only {len(devices)} local " + "JAX devices are available." + ) + if n_live % n_devices: + raise ValueError(f"n_live={n_live} must be divisible by n_devices={n_devices}.") + if n_delete % n_devices: + raise ValueError( + f"n_delete={n_delete} must be divisible by n_devices={n_devices}." + ) + + return Mesh(np.asarray(devices[:n_devices]), axis_names=(_LIVE_AXIS,)) + + +def place_state(state, mesh: Mesh): + """Shard live particles and replicate adaptive/integrator state.""" + live = NamedSharding(mesh, P(_LIVE_AXIS)) + replicated = NamedSharding(mesh, P()) + return state._replace( + particles=jax.tree.map(lambda x: jax.device_put(x, live), state.particles), + inner_kernel_params=jax.tree.map( + lambda x: jax.device_put(x, replicated), state.inner_kernel_params + ), + integrator=jax.tree.map( + lambda x: jax.device_put(x, replicated), state.integrator + ), + ) + + +def place_key(rng_key, mesh: Mesh): + """Replicate a sampler PRNG key over the mesh.""" + return jax.device_put(rng_key, NamedSharding(mesh, P())) + + +def _all_gather(mesh: Mesh, value, in_spec): + def gather(x): + return jax.lax.all_gather(x, _LIVE_AXIS, tiled=True) + + return jax.shard_map( + gather, + mesh=mesh, + in_specs=in_spec, + out_specs=P(), + check_vma=False, + )(value) + + +def delete_fn_sharded(mesh: Mesh, n_delete: int) -> Callable: + """Select the globally lowest-likelihood live particles.""" + + def delete_fn(state): + global_log_likelihood = _all_gather( + mesh, state.particles.loglikelihood, P(_LIVE_AXIS) + ) + _, dead_idx = jax.lax.top_k(-global_log_likelihood, n_delete) + return dead_idx, dead_idx + + return delete_fn + + +def 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.""" + live = NamedSharding(mesh, P(_LIVE_AXIS)) + + def update_function(rng_key, state, loglikelihood_0, **step_parameters): + particles = state.particles + particle_specs = jax.tree.map( + lambda x: P(_LIVE_AXIS) if x.ndim > 0 else P(), particles + ) + global_particles = jax.shard_map( + lambda xs: jax.tree.map( + lambda x: jax.lax.all_gather(x, _LIVE_AXIS, tiled=True), xs + ), + mesh=mesh, + in_specs=(particle_specs,), + out_specs=jax.tree.map(lambda _: P(), particles), + check_vma=False, + )(particles) + global_log_likelihood = global_particles.loglikelihood + + survivor_weights = (global_log_likelihood > loglikelihood_0).astype(jnp.float32) + survivor_weights = jnp.where( + survivor_weights.sum() > 0, + survivor_weights, + jnp.ones_like(survivor_weights), + ) + + choice_key, sample_key = jax.random.split(rng_key) + start_idx = jax.random.choice( + choice_key, + global_log_likelihood.shape[0], + shape=(n_delete,), + p=survivor_weights / survivor_weights.sum(), + replace=True, + ) + start_states = jax.tree.map(lambda x: x[start_idx], global_particles) + start_states = jax.tree.map(lambda x: jax.device_put(x, live), start_states) + sample_keys = jax.device_put(jax.random.split(sample_key, n_delete), live) + + state_specs = jax.tree.map( + lambda x: P(_LIVE_AXIS) if x.ndim > 0 else P(), start_states + ) + parameter_specs = jax.tree.map(lambda _: P(), step_parameters) + + def run_local_chains(keys, states, threshold, parameters): + def run_chain(key, chain_state): + keys = jax.random.split(key, n_inner_steps) + + def body_fn(current_state, step_key): + return constrained_step_fn( + step_key, + current_state, + threshold, + **parameters, + ) + + return jax.lax.scan(body_fn, chain_state, keys) + + return jax.vmap(run_chain)(keys, states) + + return jax.shard_map( + run_local_chains, + mesh=mesh, + in_specs=(P(_LIVE_AXIS), state_specs, P(), parameter_specs), + out_specs=(state_specs, P(_LIVE_AXIS)), + check_vma=False, + )(sample_keys, start_states, loglikelihood_0, step_parameters) + + return update_function + + +def 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.""" + inner_kernel = update_with_mcmc_take_last_sharded( + constrained_step_fn, + n_inner_steps=n_inner_steps, + n_delete=n_delete, + mesh=mesh, + ) + return build_adaptive_kernel( + delete_fn_sharded(mesh, n_delete), + inner_kernel, + update_inner_kernel_params_fn=update_inner_kernel_params_fn, + ) diff --git a/src/jimgw/samplers/blackjax/swig.py b/src/jimgw/samplers/blackjax/swig.py index e1424c8a..420866d2 100644 --- a/src/jimgw/samplers/blackjax/swig.py +++ b/src/jimgw/samplers/blackjax/swig.py @@ -15,6 +15,7 @@ from blackjax.smc.tuning.from_particles import particles_covariance_matrix from jimgw.samplers.blackjax.nss import BlackJAXNSSSampler +from jimgw.samplers.blackjax.sharding import build_sharded_from_mcmc_kernel from jimgw.samplers.config import BlackJAXSwiGConfig from jimgw.samplers.periodic import _build_masks_arrays @@ -214,7 +215,7 @@ def _checkpoint_tag(self) -> str: def _update_inner_kernel_params_fn(self) -> Callable: return self._block_covariance_update - def _build_nested_sampler(self, n_delete: int) -> SamplingAlgorithm: + def _build_nested_sampler(self, n_delete: int, mesh=None) -> SamplingAlgorithm: constrained_step = _build_swig_constrained_step( log_prior_fn=self._log_prior_fn, build_cache_fn=self._build_cache_fn, @@ -228,10 +229,19 @@ def _build_nested_sampler(self, n_delete: int) -> SamplingAlgorithm: periodic=self._periodic, n_dims=self.n_dims, ) - kernel = build_from_mcmc_kernel( - constrained_step, - num_inner_steps=1, - update_inner_kernel_params_fn=self._block_covariance_update, - num_delete=n_delete, - ) + if mesh is None: + kernel = build_from_mcmc_kernel( + constrained_step, + num_inner_steps=1, + update_inner_kernel_params_fn=self._block_covariance_update, + num_delete=n_delete, + ) + else: + kernel = build_sharded_from_mcmc_kernel( + constrained_step, + n_inner_steps=1, + update_inner_kernel_params_fn=self._block_covariance_update, + n_delete=n_delete, + mesh=mesh, + ) return SamplingAlgorithm(lambda position, rng_key=None: position, kernel) diff --git a/src/jimgw/samplers/config.py b/src/jimgw/samplers/config.py index b68c5458..7d23224e 100644 --- a/src/jimgw/samplers/config.py +++ b/src/jimgw/samplers/config.py @@ -117,6 +117,19 @@ def configure_jax_cache(self) -> None: jax.config.update("jax_persistent_cache_min_compile_time_secs", 0.0) +class _ShardingMixin(BaseModel): + """Single-host device sharding shared by BlackJAX nested samplers.""" + + n_devices: int = 1 + + @field_validator("n_devices") + @classmethod + def _positive_n_devices(cls, value: int) -> int: + if value < 1: + raise ValueError("n_devices must be >= 1") + return value + + # --------------------------------------------------------------------------- # flowMC sub-configs # --------------------------------------------------------------------------- @@ -316,7 +329,7 @@ def _n_live_n_delete_consistency(self) -> Self: return self -class BlackJAXNSSConfig(BaseSamplerConfig, _CheckpointMixin): +class BlackJAXNSSConfig(BaseSamplerConfig, _CheckpointMixin, _ShardingMixin): """Configuration for the BlackJAX nested slice sampler. !!! note @@ -349,10 +362,14 @@ def _n_live_n_delete_consistency(self) -> Self: f"yields n_delete = {n_delete}; require n_delete >= 1. " "Increase n_live or n_delete_frac." ) + if self.n_devices > 1 and self.n_live % self.n_devices: + raise ValueError("n_live must be divisible by n_devices when sharding") + if self.n_devices > 1 and n_delete % self.n_devices: + raise ValueError("n_delete must be divisible by n_devices when sharding") return self -class BlackJAXSwiGConfig(BaseSamplerConfig, _CheckpointMixin): +class BlackJAXSwiGConfig(BaseSamplerConfig, _CheckpointMixin, _ShardingMixin): """Configuration for Nested Slice within Gibbs (SwiG) sampling. ``blocks`` are expressed in Jim's sampling-space parameter names. Jim @@ -414,6 +431,10 @@ def _n_live_n_delete_consistency(self) -> Self: f"n_live * n_delete_frac = {self.n_live * self.n_delete_frac} " f"yields n_delete = {n_delete}; require n_delete >= 1." ) + if self.n_devices > 1 and self.n_live % self.n_devices: + raise ValueError("n_live must be divisible by n_devices when sharding") + if self.n_devices > 1 and n_delete % self.n_devices: + raise ValueError("n_delete must be divisible by n_devices when sharding") return self diff --git a/tests/unit/samplers/blackjax/test_sharding.py b/tests/unit/samplers/blackjax/test_sharding.py new file mode 100644 index 00000000..8ef0dfff --- /dev/null +++ b/tests/unit/samplers/blackjax/test_sharding.py @@ -0,0 +1,154 @@ +"""Multi-device tests for BlackJAX NSS and SwiG.""" + +from __future__ import annotations + +import pickle +from pathlib import Path + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jimgw.samplers.blackjax.nss import BlackJAXNSSSampler +from jimgw.samplers.blackjax.sharding import make_live_mesh +from jimgw.samplers.blackjax.swig import BlackJAXSwiGSampler +from jimgw.samplers.config import BlackJAXNSSConfig, BlackJAXSwiGConfig + +_HAS_FOUR_DEVICES = jax.local_device_count() >= 4 + + +def _log_prior(position): + return jnp.where(jnp.all((position >= 0.0) & (position <= 1.0)), 0.0, -jnp.inf) + + +def _build_cache(position): + return position[0] ** 2 + + +def _log_likelihood_from_cache(position, cache): + return -20.0 * ((cache - 0.25) ** 2 + (position[1] - 0.5) ** 2) + + +def _log_likelihood(position): + return _log_likelihood_from_cache(position, _build_cache(position)) + + +def test_make_live_mesh_rejects_unavailable_devices(): + with pytest.raises(ValueError, match="only .* local JAX devices"): + make_live_mesh(jax.local_device_count() + 1, 16, 4) + + +@pytest.mark.skipif( + not _HAS_FOUR_DEVICES, + reason="run with XLA_FLAGS=--xla_force_host_platform_device_count=4", +) +def test_nss_runs_sharded_and_preserves_diagnostics(): + config = BlackJAXNSSConfig( + n_live=16, + n_delete_frac=0.25, + num_inner_steps_per_dim=1, + termination_dlogz=2.0, + n_devices=4, + ) + sampler = BlackJAXNSSSampler( + n_dims=2, + log_prior_fn=_log_prior, + log_likelihood_fn=_log_likelihood, + log_posterior_fn=lambda x: _log_prior(x) + _log_likelihood(x), + config=config, + ) + + initial = jax.random.uniform(jax.random.key(0), (16, 2)) + sampler.sample(jax.random.key(1), initial) + diagnostics = sampler.get_diagnostics() + + assert diagnostics["n_iterations"] > 0 + assert diagnostics["n_likelihood_evaluations"] > 0 + stepping_out = diagnostics["n_stepping_out_history"] + assert stepping_out.shape[-1] == 2 + assert stepping_out.shape[0] % config.n_devices == 0 + + +@pytest.mark.skipif( + not _HAS_FOUR_DEVICES, + reason="run with XLA_FLAGS=--xla_force_host_platform_device_count=4", +) +def test_swig_runs_sharded_with_consistent_cache(): + config = BlackJAXSwiGConfig( + blocks=[["slow"], ["fast"]], + n_live=16, + n_delete_frac=0.25, + num_gibbs_sweeps=1, + max_steps=3, + max_shrinkage=20, + termination_dlogz=2.0, + n_devices=4, + ) + sampler = BlackJAXSwiGSampler( + n_dims=2, + log_prior_fn=_log_prior, + log_likelihood_fn=_log_likelihood, + log_posterior_fn=lambda x: _log_prior(x) + _log_likelihood(x), + config=config, + block_indices=((0,), (1,)), + refresh_cache=(True, False), + build_cache_fn=_build_cache, + log_likelihood_from_cache_fn=_log_likelihood_from_cache, + ) + + initial = jax.random.uniform(jax.random.key(2), (16, 2)) + sampler.sample(jax.random.key(3), initial) + result = sampler.get_samples() + expected = jax.vmap(_log_likelihood)(jnp.asarray(result["samples"])) + + np.testing.assert_allclose(result["log_likelihood"], expected, rtol=2e-6) + assert sampler.get_diagnostics()["n_likelihood_evaluations"] > 0 + assert not hasattr(sampler._final_state.particles, "cache") + + +@pytest.mark.skipif( + not _HAS_FOUR_DEVICES, + reason="run with XLA_FLAGS=--xla_force_host_platform_device_count=4", +) +def test_sharded_checkpoint_is_host_backed_and_resumable(tmp_path, monkeypatch): + config = BlackJAXNSSConfig( + n_live=16, + n_delete_frac=0.25, + num_inner_steps_per_dim=1, + termination_dlogz=2.0, + n_devices=4, + checkpoint_dir=tmp_path, + checkpoint_interval=1e-9, + ) + + def make_sampler(): + return BlackJAXNSSSampler( + n_dims=2, + log_prior_fn=_log_prior, + log_likelihood_fn=_log_likelihood, + log_posterior_fn=lambda x: _log_prior(x) + _log_likelihood(x), + config=config, + ) + + checkpoint = tmp_path / "checkpoint.pkl" + original_unlink = Path.unlink + monkeypatch.setattr( + Path, + "unlink", + lambda self, missing_ok=False: ( + None if self == checkpoint else original_unlink(self, missing_ok=missing_ok) + ), + ) + initial = jax.random.uniform(jax.random.key(4), (16, 2)) + make_sampler().sample(jax.random.key(5), initial) + monkeypatch.setattr(Path, "unlink", original_unlink) + + with checkpoint.open("rb") as stream: + saved = pickle.load(stream) + assert isinstance(saved["state"].particles.position, np.ndarray) + + resumed = make_sampler() + resumed.sample(jax.random.key(999), initial) + assert resumed.get_diagnostics()["n_iterations"] > 0 + assert not checkpoint.exists() diff --git a/tests/unit/samplers/test_config.py b/tests/unit/samplers/test_config.py index 38d4918a..98727f15 100644 --- a/tests/unit/samplers/test_config.py +++ b/tests/unit/samplers/test_config.py @@ -48,9 +48,7 @@ def test_sampler_config_union_from_dict(): cfg4 = ta.validate_python({"type": "blackjax-smc"}) assert isinstance(cfg4, BlackJAXSMCConfig) - cfg5 = ta.validate_python( - {"type": "blackjax-swig", "blocks": [["x"], ["y"]]} - ) + cfg5 = ta.validate_python({"type": "blackjax-swig", "blocks": [["x"], ["y"]]}) assert isinstance(cfg5, BlackJAXSwiGConfig) @@ -67,6 +65,23 @@ def test_swig_sampling_defaults(): config = BlackJAXSwiGConfig(blocks=[["x"]]) assert config.num_gibbs_sweeps == 2 assert config.termination_dlogz == pytest.approx(math.exp(-3.0)) + assert config.n_devices == 1 + + +@pytest.mark.parametrize("config_cls", [BlackJAXNSSConfig, BlackJAXSwiGConfig]) +def test_nested_sampler_sharding_requires_divisible_particle_counts(config_cls): + kwargs = {"blocks": [["x"]]} if config_cls is BlackJAXSwiGConfig else {} + with pytest.raises(ValueError, match="n_live must be divisible by n_devices"): + config_cls(n_live=10, n_delete_frac=0.4, n_devices=4, **kwargs) + with pytest.raises(ValueError, match="n_delete must be divisible by n_devices"): + config_cls(n_live=12, n_delete_frac=0.25, n_devices=2, **kwargs) + + +@pytest.mark.parametrize("config_cls", [BlackJAXNSSConfig, BlackJAXSwiGConfig]) +def test_nested_sampler_sharding_requires_positive_device_count(config_cls): + kwargs = {"blocks": [["x"]]} if config_cls is BlackJAXSwiGConfig else {} + with pytest.raises(ValueError, match="n_devices must be >= 1"): + config_cls(n_devices=0, **kwargs) def test_extra_fields_forbidden():