2 ways to batch inference

batching.py

import jax


""" 'Loop over vmaps'

This is what most PPLs do. We synchronize chains at the step level, each step
ends when the `num_chains` have completed their step. We iterate for `num_samples`
steps.
"""

def one_step(rng_key, states):
    chain_keys = jax.random.split(rng_key, num_chains)
    new_states, _ = jax.vmap(kernel)(chain_keys, states)
    return new_states, states

step_keys = jax.random.split(rng_key, num_samples)
_, states = jax.lax.scan(one_step, initial_states, step_keys)

""" 'vmap the loops'

We run `num_chains` independent chains. This ends when all chains
have run `num_steps`.
"""

def chain(rng_key, initial_state):
    """Moves one chain `num_steps` time."""

    def one_step(state, key):
        new_state, _ = kernel(key, state)
        return new_state, state

    step_keys = jax.random.split(rng_key, num_samples)
    _, states = jax.lax.scan(one_step, initial_state, step_keys)
    return states

chain_keys = jax.random.split(rng_key, num_chains)
run = jax.vmap(chain)(keys, initial_states)