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,
    use_detector_reference_sample: bool = False,
    use_observable_reference_sample: 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. Defaults to
            None, which automatically chooses a batch size based on available
            memory. When using a GPU, setting this explicitly can help fully
            utilize VRAM for maximum performance. NOTE: Changing the batch size
            will affect reproducibility even with a fixed seed.
        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.
        use_detector_reference_sample: Defaults to False. When True, a noiseless
            reference sample is computed and XORed with detector outcomes so that
            results represent deviations from the noiseless baseline. This should
            only be used when detectors are deterministic. Otherwise, it can
            unpredictably change the results.
        use_observable_reference_sample: Defaults to False. When True, a noiseless
            reference sample is computed and XORed with observable outcomes so that
            results represent deviations from the noiseless baseline. This should
            only be used when observables are deterministic. Otherwise, it can
            unpredictably change the results.

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

    Raises:
        ValueError: If ``separate_observables`` is combined with
            ``prepend_observables`` or ``append_observables``.

    """
    if separate_observables and (prepend_observables or append_observables):
        raise ValueError(
            "Can't specify separate_observables=True with "
            "append_observables=True or prepend_observables=True"
        )

    compute_reference = (
        use_detector_reference_sample or use_observable_reference_sample
    )

    if compute_reference:
        samples, reference = self._sample_batches(
            shots, batch_size, compute_reference=True
        )
        num_detectors = self._num_detectors
        if use_detector_reference_sample:
            samples[:, :num_detectors] ^= reference[:num_detectors]
        if use_observable_reference_sample:
            samples[:, num_detectors:] ^= reference[num_detectors:]
    else:
        samples = self._sample_batches(shots, batch_size)

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

    if prepend_observables and append_observables:
        combined = np.concatenate([obs_samples, det_samples, obs_samples], axis=1)
        return _maybe_bit_pack(combined, bit_packed=bit_packed)
    if append_observables:
        return _maybe_bit_pack(samples, bit_packed=bit_packed)
    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 | None = None) -> 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. Defaults to
            None, which automatically chooses a batch size based on available
            memory. When using a GPU, setting this explicitly can help fully
            utilize VRAM for maximum performance. NOTE: Changing the batch size
            will affect reproducibility even with a fixed seed.

    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.

    """
    if batch_size < 1:
        raise ValueError(f"batch_size must be at least 1, got {batch_size}")
    expected_outputs = self._program.num_outputs
    if state.shape != (expected_outputs,):
        raise ValueError(
            f"state must have shape ({expected_outputs},), got {state.shape}"
        )
    f_samples = jnp.asarray(self._channel_sampler.sample(batch_size))
    p_norm = jnp.ones(batch_size)
    p_joint = jnp.ones(batch_size)

    if len(self._program.direct_f_indices) > 0:
        direct_bits = (
            f_samples[:, self._program.direct_f_indices].astype(jnp.bool_)
            ^ self._program.direct_flips
        )
        n_direct = len(self._program.direct_f_indices)
        targets = state[self._program.output_order[:n_direct]]
        p_joint = p_joint * (direct_bits == targets).all(axis=1)

    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)