docs: 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