Module gabenet.mcmc

Functions

def sample_markov_chain(key, kernel: Union[Callable[[tuple], tuple[Any, dict]], haiku._src.transform.TransformedWithState], n_samples: int, n_burnin_steps: int, n_chains: Optional[int] = None, params: Optional[Mapping[str, Mapping[str, jax.Array]]] = None, initial_state: Optional[Mapping[str, Mapping[str, jax.Array]]] = None, n_leap_size: int = 1)

Take samples using Markov chain Monte Carlo.

Args

key

Pseudo random number generator key.

kernel

The transition kernel that advances the Markov chain by one step. The kernel is either a haiku TransformedWithState or a function compatible with TransformedWithState.apply.

n_samples

Number of samples to generate.

n_burnin_steps

Advance the chain by this many steps before taking a sample.

n_chains

Number of (independent) Markov chains to run.

params

The chain-specific but state-independent params (a pytree) of the Markov chain. When None, kernel is assumed to be a TransformedWithState to initialise using kernel.init.

initial_state

Starting state (a pytree) of the Markov chain. Leading dimension corresponds to a chain.

n_leap_size

Number of successive steps between samples.

Returns

A pair of (params, state), each a pytree. The leading dimension of params and state refers to the chain. The second dimension of state corresponds to the samples in per chain.

Expand source code

def sample_markov_chain(
    key,
    kernel: Callable[[tuple], tuple[Any, dict]] | hk.TransformedWithState,
    n_samples: int,
    n_burnin_steps: int,
    n_chains: Optional[int] = None,
    params: Optional[hk.Params] = None,
    initial_state: Optional[hk.State] = None,
    n_leap_size: int = 1,
):
    """Take samples using Markov chain Monte Carlo.

    Args:
        key: Pseudo random number generator key.
        kernel: The transition kernel that advances the Markov chain by one step. The
            kernel is either a haiku `TransformedWithState` or a function compatible
            with `TransformedWithState.apply`.
        n_samples: Number of samples to generate.
        n_burnin_steps: Advance the chain by this many steps before taking a sample.
        n_chains: Number of (independent) Markov chains to run.
        params: The chain-specific but state-independent `params` (a pytree) of
            the Markov chain. When `None`, kernel is assumed to be a
            `TransformedWithState` to initialise using `kernel.init`.
        initial_state: Starting `state` (a pytree) of the Markov chain. Leading
            dimension corresponds to a chain.
        n_leap_size: Number of successive steps between samples.

    Returns:
        A pair of `(params, state)`, each a pytree. The leading dimension of
        `params` and `state` refers to the chain. The second dimension of `state`
        corresponds to the samples in per chain.
    """
    key_seq = hk.PRNGSequence(key)

    _scannable_kernel_fn = scannable(kernel)

    if initial_state is None or params is None:
        if not isinstance(kernel, hk.TransformedWithState):
            raise ValueError(
                "Kernel must be `TransformedWithState` to be able to initialise"
                "`params` and `state` when either is None."
            )
        if n_chains is None:
            raise ValueError("Number of chains not specified!")

        init_key_per_chain = random.split(next(key_seq), num=n_chains)
        params, initial_state = jax.vmap(kernel.init)(init_key_per_chain)
    else:
        states_leaves = jax.tree_util.tree_leaves(initial_state)
        # Leading axis of states corresponds to chains.
        n_chains_expected = states_leaves[0].shape[0]
        if n_chains is None:
            n_chains = n_chains_expected
        elif n_chains != n_chains_expected:
            raise ValueError(
                f"The leading axis of `initial_state` suggests "
                f"{n_chains_expected} chains, but the argument is set to "
                f"`n_chains={n_chains}`."
            )

    def _leapfrog(carry, n_steps: int):
        """A for-loop (from 0,..,`n_steps-1`) that runs `step`."""
        carry_out, _ = jax.lax.scan(_scannable_kernel_fn, carry, None, length=n_steps)
        state = carry_out[1]
        return carry_out, state

    def _sampler(params_init, state_init, key, n_sample_size: int):
        """Samper for a single Markov chain."""
        carry = (params_init, state_init, key)
        # 1) Take first sample after `(n_burnin - n_leap_size) + n_leap_size` steps.
        # 2) Take subsequent samples after `n_leap_size` steps.
        carry, _ = _leapfrog(carry, n_steps=n_burnin_steps - n_leap_size)
        carry, stacked_states = jax.lax.scan(
            lambda c, _: _leapfrog(c, n_leap_size),
            carry,
            xs=None,
            length=n_sample_size,
        )

        return stacked_states

    n_devices = jax.local_device_count()

    if n_devices == 1:
        warnings.warn("Only one visible device in JAX. Reconfigure XLA_FLAGS.")

    if n_devices == n_chains:
        # Run each chain on a separate device.
        _vectorised_kernel_fn = jax.pmap(
            _sampler, in_axes=(0, 0, 0, None), static_broadcasted_argnums=3
        )
    else:
        # Vectorise with vmap instead of sharding across devices.
        warnings.warn(
            f"Chains (n={n_chains}) not divisible across devices. Falling back to vmap."
        )
        _vectorised_kernel_fn = jax.vmap(_sampler, in_axes=(0, 0, 0, None))

    if n_samples % n_chains > 0:
        raise ValueError(
            "Number of samples {n_samples} not divisible by number of chains {n_chains}."
        )

    # Generate an initial state for each chain.
    keys = random.split(next(key_seq), num=n_chains)

    n_samples_per_chain = n_samples // n_chains
    state = _vectorised_kernel_fn(params, initial_state, keys, n_samples_per_chain)
    return params, state