def _probs(
    density_matrix: np.ndarray, indices: Sequence[int], qid_shape: tuple[int, ...]
) -> np.ndarray:
    """Returns the probabilities for a measurement on the given indices."""
    # Only diagonal elements matter.
    all_probs = np.diagonal(np.reshape(density_matrix, (np.prod(qid_shape, dtype=np.int64),) * 2))

    return simulation_utils.state_probabilities_by_indices(all_probs.real, indices, qid_shape)
Returns the probabilities for a measurement on the given indices.
cirq-core/cirq/sim/density_matrix_utils.py