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,
    postselection_mask: np.ndarray | None = None,
) -> 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 at once. Defaults to None,
            which chooses it from the available memory. With
            ``strategy="active_statevector"`` it bounds the shots a thread
            holds and does not change the samples.
            With a stabilizer-rank strategy on a GPU, setting this explicitly
            can help fully utilize VRAM; it changes the samples of a 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. The
            packed rows may be a strided view of a wider array.
        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.
        postselection_mask: Optional boolean array of length ``num_detectors``.
            When set, shots where any masked direct detector (an error bit
            or a constant) fires are not sampled further. All ``shots`` rows
            are still returned: survivors
            contain the full sample, while discarded rows retain direct detector
            columns and fill component columns with ``False``; with
            ``strategy="active_statevector"``, a shot discarded by a masked
            component detector also reads ``False`` in the outputs set after
            it. Re-apply the mask to the returned detector columns to recover
            surviving shots.

    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
    )
    num_detectors = self._num_detectors

    if postselection_mask is not None:
        mask = np.asarray(postselection_mask, dtype=np.bool_)
        if mask.shape != (num_detectors,):
            raise ValueError(
                f"postselection_mask must have shape ({num_detectors},), "
                f"got {mask.shape}"
            )
        # NOTE: shot words discard and prune with any mask; with a
        # stabilizer-rank strategy only direct detectors that can fire skip work.
        if self._shots is not None:
            postselection_mask = mask if self._program.components else None
        else:
            skips = mask & self._direct_detector_mask
            indices, flips, columns = self._direct_detectors
            constant = indices < 0
            skips[columns[constant]] &= flips[constant] & (
                not use_detector_reference_sample
            )
            postselection_mask = mask if skips.any() else None

    num_outputs = self._program.num_outputs
    is_detector = np.arange(num_outputs) < num_detectors
    referenced = np.where(
        is_detector, use_detector_reference_sample, use_observable_reference_sample
    )
    columns = _output_columns(
        num_detectors,
        num_outputs,
        prepend=prepend_observables,
        append=append_observables,
        separate=separate_observables,
    )
    if self._shots is not None:
        # NOTE: the reference bits are XOR-ed in before the postselection tests.
        flip = self._reference_sample() & referenced if compute_reference else None
        if self._program.components or bit_packed:
            arrays = self._sample_shot_outputs(
                shots,
                batch_size,
                columns,
                bit_packed=bit_packed,
                mask=postselection_mask,
                flip=flip,
            )
            return (arrays[0], arrays[1]) if separate_observables else arrays[0]
        _check_sizes(shots, batch_size)
        samples = self._shots.direct_bools(shots, flip)
    elif postselection_mask is not None:
        if not compute_reference:
            referenced = None
        samples = self._sample_batches_with_postselection(
            shots,
            batch_size,
            postselection_mask=postselection_mask,
            referenced=referenced,
        )
    elif compute_reference:
        samples, reference = self._sample_batches(
            shots, batch_size, compute_reference=True
        )
        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)
    arrays = [
        _maybe_bit_pack(_take_columns(samples, c), bit_packed=bit_packed)
        for c in columns
    ]
    return (arrays[0], arrays[1]) if separate_observables else arrays[0]

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 at once. Defaults to None,
            which chooses it from the available memory. With
            ``strategy="active_statevector"`` it bounds the shots a thread
            holds and does not change the samples.
            With a stabilizer-rank strategy on a GPU, setting this explicitly
            can help fully utilize VRAM; it changes the samples of a seed.

    Returns:
        A numpy array containing the measurement samples.

    """
    if self._shots is not None:
        columns = [np.arange(self._program.num_outputs)]
        if not self._program.components:
            _check_sizes(shots, batch_size)
            return self._shots.direct_bools(shots)
        return self._sample_shot_outputs(
            shots, batch_size, columns, bit_packed=False
        )[0]
    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 = _direct_bits(self._program, f_samples)
        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 isinstance(component, CompiledComponent)
        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)