sequince.steps.mutation.blackjax
================================

.. py:module:: sequince.steps.mutation.blackjax


Attributes
----------

.. autoapisummary::

   sequince.steps.mutation.blackjax.logger
   sequince.steps.mutation.blackjax.BlackJAXAlgorithmFactory


Classes
-------

.. autoapisummary::

   sequince.steps.mutation.blackjax.BlackJAXOutput
   sequince.steps.mutation.blackjax.BlackJAXMutation


Module Contents
---------------

.. py:data:: logger

.. py:data:: BlackJAXAlgorithmFactory

.. py:class:: BlackJAXOutput

   Bases: :py:obj:`collections.abc.Mapping`\ [\ :py:obj:`str`\ , :py:obj:`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.


   .. py:attribute:: mean_acceptance_probability
      :type:  Any


   .. py:attribute:: divergence
      :type:  Any


   .. py:attribute:: algorithm
      :type:  str


   .. py:attribute:: steps
      :type:  int


   .. py:attribute:: chain
      :type:  Any


   .. py:attribute:: num_integration_steps
      :type:  Any
      :value: None



   .. py:attribute:: metrics
      :type:  Any
      :value: None



.. py:class:: BlackJAXMutation(log_prob_fn, *, n_steps = 1, algorithm = 'nuts', algorithm_kwargs = None, transform = None, transform_refit = 'always', store_chains = False, metrics = ())

   Bases: :py:obj:`sequince.steps.mutation.base.MutationStep`


   Run independent BlackJAX chains to rejuvenate a JAX particle cloud.

   :param log_prob_fn: JAX-compatible scalar log-probability function or a
                       :class:`~sequince.target.Target` whose component functions
                       are JAX compatible.
   :param n_steps: Number of BlackJAX transitions performed for every particle.
   :param algorithm: 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``.
   :param algorithm_kwargs: 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``.
   :param transform: Optional particle transform applied before mutation.
   :param transform_refit: Policy controlling when the particle transform is refitted.
   :param store_chains: Whether to include all BlackJAX positions in the step output.

   .. rubric:: 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.


   .. py:property:: algorithm_name
      :type: str


      Human-readable name of the configured BlackJAX algorithm.


   .. py:method:: apply(state, context, rng)


   .. py:method:: apply_with_mutation_state(state, context, rng, *, n_steps, mutation_state = None)

      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.