def _probs(density_matrix: np.ndarray, indices: List[int],
    num_qubits: int) -> List[float]:
    """Returns the probabilities for a measurement on the given indices."""
    # Only diagonal elements matter.
    all_probs = np.diagonal(
        np.reshape(density_matrix, (2 ** num_qubits, 2 ** num_qubits)))
    # Shape into a tensor
    tensor = np.reshape(all_probs, [2] * num_qubits)

    # Calculate the probabilities for measuring the particular results.
    probs = [
        np.sum(np.abs(tensor[linalg.slice_for_qubits_equal_to(indices, b)]))
        for b in range(2 ** len(indices))]

    # To deal with rounding issues, ensure that the probabilities sum to 1.
    probs /= np.sum(probs) # type: ignore
    return probs
Returns the probabilities for a measurement on the given indices.