execute

def choose_outcome(p0: Array, p1: Array, count: Array, key: Array) -> Array:
    """Return each lane's outcome: True with probability ``p1 = 1 - p0``.

    The less likely outcome is drawn, so that both outcomes keep the relative
    precision of their probability: in float32, ``p1`` alone rounds to 1 once
    ``p0`` is below about 6·10⁻⁸. Two uint32 draws per lane form a uniform 64-bit
    fraction, compared with ``p·2^32`` (an exact scaling) for the smaller ``p``,
    so the resolution is 2^-64 for any float dtype, and an outcome of
    probability 0 is never drawn. ``count`` is unused while every lane is one
    shot.
    """
    del count
    high, low = jax.random.bits(key, (2, *p1.shape), dtype=jnp.uint32)
    one_is_rare = p1 <= p0
    scaled = jnp.clip(jnp.minimum(p0, p1), 0, 0.5) * _TWO_TO_32
    whole = jnp.floor(scaled)
    high_bound = whole.astype(jnp.uint32)
    low_bound = jnp.floor((scaled - whole) * _TWO_TO_32).astype(jnp.uint32)
    rare = (high < high_bound) | ((high == high_bound) & (low < low_bound))
    return rare == one_is_rare

def initial_state(component: StatevectorComponent, f_bits: Array) -> LaneState:
    """Return lanes in ``|0…0⟩`` with the given f-bits, shape ``(L, num_f)``."""
    lanes = f_bits.shape[0]
    dtype = jnp.dtype(component.dtype)
    amps = jnp.zeros((lanes, 2**component.width), dtype=dtype).at[:, 0].set(1)
    return LaneState(
        amps=amps,
        symbols=pack_bits(f_bits, component.num_words),
        count=jnp.ones(lanes, dtype=jnp.int32),
        max_norm_deviation=jnp.zeros((), dtype=jnp.finfo(dtype).dtype),
    )

def lane_bytes(component: StatevectorComponent) -> int:
    """Return the peak device memory of one lane, temporary amplitudes included."""
    # NOTE: XLA's memory analysis of the executor shows four amplitude arrays.
    amplitudes = 4 * 2**component.width * jnp.dtype(component.dtype).itemsize
    return amplitudes + 4 * component.num_words + component.num_outputs

def outputs(component: StatevectorComponent, state: LaneState) -> Array:
    """Return each lane's output bits, shape ``(L, num_outputs)``."""
    return _affine(
        state.symbols[:, None, :], component.output_row, component.output_const
    )

def pack_bits(bits: Array, num_words: int) -> Array:
    """Pack boolean columns into uint32 words: column s is bit s % 32 of word s // 32."""
    lanes, num_bits = bits.shape
    padded = jnp.zeros((lanes, num_words * 32), dtype=jnp.uint32)
    padded = padded.at[:, :num_bits].set(bits.astype(jnp.uint32))
    shifts = jnp.arange(32, dtype=jnp.uint32)
    words = padded.reshape(lanes, num_words, 32) << shifts
    return jnp.sum(words, axis=-1, dtype=jnp.uint32)

def run(
    component: StatevectorComponent,
    state: LaneState,
    key: Array,
    start: int | Array,
    stop: int | Array,
) -> LaneState:
    """Run actions ``start … stop-1`` of the tape on every lane.

    A component without actions returns ``state`` unchanged.
    """
    if component.num_actions == 0:
        return state
    branches = [partial(_KERNELS[kind], component) for kind in range(len(_KERNELS))]

    def step(a: Array, carry: LaneState) -> LaneState:
        return jax.lax.switch(component.kind[a], branches, carry, a, key)

    return jax.lax.fori_loop(start, stop, step, state)

@jax.jit
def sample_statevector_component(
    component: StatevectorComponent, f_params: Array, key: Array
) -> tuple[Array, Array, Array]:
    """Sample one component's outputs; the contract of ``sampler.sample_component``.

    Args:
        component: The compiled component.
        f_params: Error-basis bits of every shot, shape ``(B, num_f_global)``.
        key: JAX random key.

    Returns:
        ``(samples, next_key, max_norm_deviation)`` with ``samples`` of shape
        ``(B, num_outputs)``.

    """
    f_bits = f_params[:, component.f_selection].astype(jnp.bool_)
    key, subkey = jax.random.split(key)
    state = initial_state(component, f_bits)
    state = run(component, state, subkey, 0, component.num_actions)
    return outputs(component, state), key, state.max_norm_deviation