Get the probabilities of measuring a qubit.

    :param wf: The statevector with a dimension for each qubit
    :param qubit: The qubit to measure. We will sum over every axis except this one.
    :return: A vector of classical probabilities.

def get_measure_probabilities(wf, qubit):
    """
    Get the probabilities of measuring a qubit.

    :param wf: The statevector with a dimension for each qubit
    :param qubit: The qubit to measure. We will sum over every axis except this one.
    :return: A vector of classical probabilities.
    """
    n_qubits = len(wf.shape)
    all_inds = list(range(n_qubits))

    return np.einsum(np.conj(wf), all_inds,
                     wf, all_inds,
                     [int(qubit)])