sampler

def sample(
    self,
    shots: int,
    *,
    batch_size: int | None = None,
    prepend_observables: bool = False,
    append_observables: bool = False,
    separate_observables: bool = False,
    bit_packed: bool = False,
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    """Return detector samples from the circuit.

    The circuit must define the detectors using DETECTOR instructions. Observables
    defined by OBSERVABLE_INCLUDE instructions can also be included in the results
    as honorary detectors.

    Args:
        shots: The number of times to sample every detector in the circuit.
        batch_size: The number of samples to process in each batch. When using a
            GPU, it is recommended to increase this value until VRAM is fully
            utilized for maximum performance.
        separate_observables: Defaults to False. When set to True, the return value
            is a (detection_events, observable_flips) tuple instead of a flat
            detection_events array.
        prepend_observables: Defaults to false. When set, observables are included
            with the detectors and are placed at the start of the results.
        append_observables: Defaults to false. When set, observables are included
            with the detectors and are placed at the end of the results.
        bit_packed: Defaults to false. When set, results are bit-packed.

    Returns:
        A numpy array or tuple of numpy arrays containing the samples.

    """
    samples = self._sample_batches(shots, batch_size)

    if append_observables:
        return _maybe_bit_pack(samples, bit_packed=bit_packed)

    num_detectors = self._num_detectors
    det_samples = samples[:, :num_detectors]
    obs_samples = samples[:, num_detectors:]

    if prepend_observables:
        combined = np.concatenate([obs_samples, det_samples], axis=1)
        return _maybe_bit_pack(combined, bit_packed=bit_packed)
    if separate_observables:
        return (
            _maybe_bit_pack(det_samples, bit_packed=bit_packed),
            _maybe_bit_pack(obs_samples, bit_packed=bit_packed),
        )

    return _maybe_bit_pack(det_samples, bit_packed=bit_packed)

def sample(self, shots: int, *, batch_size: int = 1024) -> np.ndarray:
    """Sample measurement outcomes from the circuit.

    Args:
        shots: The number of times to sample every measurement in the circuit.
        batch_size: The number of samples to process in each batch. When using a
            GPU, it is recommended to increase this value until VRAM is fully
            utilized for maximum performance.

    Returns:
        A numpy array containing the measurement samples.

    """
    return self._sample_batches(shots, batch_size)

def probability_of(self, state: np.ndarray, *, batch_size: int) -> np.ndarray:
    """Compute probabilities for a batch of error samples given a measurement state.

    Args:
        state: The measurement outcome state to compute probability for.
        batch_size: Number of error samples to use for estimation.

    Returns:
        Array of probabilities P(state | error_sample) for each error sample.

    """
    f_samples = jnp.asarray(self._channel_sampler.sample(batch_size))
    p_norm = jnp.ones(batch_size)
    p_joint = jnp.ones(batch_size)

    for component in self._program.components:
        assert len(component.compiled_scalar_graphs) == 2

        f_selected = f_samples[:, component.f_selection]

        norm_circuit, joint_circuit = component.compiled_scalar_graphs

        # Normalization: only f-params
        p_norm = p_norm * jnp.abs(evaluate(norm_circuit, f_selected))

        # Joint probability: f-params + state
        component_state = state[list(component.output_indices)]
        tiled_state = jnp.tile(component_state, (batch_size, 1))
        joint_params = jnp.hstack([f_selected, tiled_state])
        p_joint = p_joint * jnp.abs(evaluate(joint_circuit, joint_params))

    return np.asarray(p_joint / p_norm)