sequince.steps.mutation.blackjax

sequince.steps.mutation.blackjax#

Attributes#

logger

BlackJAXAlgorithmFactory

Classes#

BlackJAXOutput

BlackJAX results with block acceptance in metrics.acceptance.

BlackJAXMutation

Run independent BlackJAX chains to rejuvenate a JAX particle cloud.

Module Contents#

sequince.steps.mutation.blackjax.logger#
sequince.steps.mutation.blackjax.BlackJAXAlgorithmFactory#
class sequince.steps.mutation.blackjax.BlackJAXOutput[source]#

Bases: collections.abc.Mapping[str, Any]

BlackJAX results with block acceptance in metrics.acceptance.

mean_acceptance_probability averages the backend’s MH acceptance probabilities when available; it is not the realized accepted fraction.

mean_acceptance_probability: Any#
divergence: Any#
algorithm: str#
steps: int#
chain: Any#
num_integration_steps: Any = None#
metrics: Any = None#
class sequince.steps.mutation.blackjax.BlackJAXMutation(log_prob_fn, *, n_steps=1, algorithm='nuts', algorithm_kwargs=None, transform=None, transform_refit='always', store_chains=False, metrics=())[source]#

Bases: sequince.steps.mutation.base.MutationStep

Run independent BlackJAX chains to rejuvenate a JAX particle cloud.

Parameters:

  • log_prob_fn (collections.abc.Callable[..., sequince.typing.Array | float] | sequince.target.Target) – JAX-compatible scalar log-probability function or a Target whose component functions are JAX compatible.

  • n_steps (int) – Number of BlackJAX transitions performed for every particle.

  • algorithm (str | BlackJAXAlgorithmFactory) – Named algorithm ("nuts", "hmc", or "random-walk"), or a callable that accepts the scalar log-density function and algorithm_kwargs and returns a BlackJAX SamplingAlgorithm.

  • algorithm_kwargs (collections.abc.Mapping[str, Any] | None) – Options passed to the selected BlackJAX algorithm. NUTS and HMC default to a step size of 1e-3 and a unit diagonal inverse mass matrix. HMC additionally defaults to 10 integration steps. The Gaussian random walk accepts sigma and defaults to 0.1.

  • transform (sequince.transforms.ParticleTransform | None) – Optional particle transform applied before mutation.

  • transform_refit (sequince.steps.mutation.base.TransformRefit) – Policy controlling when the particle transform is refitted.

  • store_chains (bool) – Whether to include all BlackJAX positions in the step output.

  • metrics (tuple[sequince.steps.mutation.metrics.MutationMetric, ...] | list[sequince.steps.mutation.metrics.MutationMetric])

Notes

The mutation derives a fresh JAX key from Sequince’s supplied RNG on every application. Consequently checkpointing and resumption do not require a separate BlackJAX RNG state.

property algorithm_name: str#

Human-readable name of the configured BlackJAX algorithm.

Return type:

str

apply(state, context, rng)[source]#
Parameters:
Return type:

tuple[sequince.state.ParticleState, sequince.steps.context.StepContext, BlackJAXOutput]

apply_with_mutation_state(state, context, rng, *, n_steps, mutation_state=None)[source]#

Run transitions within a block without finalizing metrics.

Pass mutation_state=None to initialize a new block; otherwise pass the state returned by the preceding call. The state holds individual metric states and cumulative acceptance bookkeeping. It must belong to the same metric configuration and mutation coordinate space.

Returns (state, context, output, mutation_state). The output’s metrics remain unset, including on the first and zero-step calls. Concrete mutations prepare existing metrics on entry, update state for each transition, and return it explicitly without calling metric.result(). Runtime state must never be stored on the mutation object.

Use apply or apply_with_n_steps for a complete standalone block. RepeatedMutationStep uses this method to continue across rounds and finalizes metrics once after the complete block.

Parameters:
Return type:

tuple[sequince.state.ParticleState, sequince.steps.context.StepContext, BlackJAXOutput, sequince.steps.mutation.metrics.MutationStepState | None]