Convert between bitstring probabilities and joint Z moment expectations.

    :param np.array p: An array that enumerates bitstring probabilities. When
        flattened out ``p = [p_00...0, p_00...1, ...,p_11...1]``. The total number of elements must
        therefore be a power of 2. The canonical shape has a separate axis for each qubit, such that
        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
    :return: ``z_moments``, an np.array with one length-2 axis per qubit which contains the
        expectations of all monomials in ``{I, Z_0, Z_1, ..., Z_{n-1}}``. The expectations of each
        monomial can be accessed via::

            <Z_0^j_0 Z_1^j_1 ... Z_m^j_m> = z_moments[j_0,j_1,...,j_m]

    :rtype: np.array

def bitstring_probs_to_z_moments(p):
    """
    Convert between bitstring probabilities and joint Z moment expectations.

    :param np.array p: An array that enumerates bitstring probabilities. When
        flattened out ``p = [p_00...0, p_00...1, ...,p_11...1]``. The total number of elements must
        therefore be a power of 2. The canonical shape has a separate axis for each qubit, such that
        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
    :return: ``z_moments``, an np.array with one length-2 axis per qubit which contains the
        expectations of all monomials in ``{I, Z_0, Z_1, ..., Z_{n-1}}``. The expectations of each
        monomial can be accessed via::

            <Z_0^j_0 Z_1^j_1 ... Z_m^j_m> = z_moments[j_0,j_1,...,j_m]

    :rtype: np.array
    """
    zmat = np.array([[1, 1],
                     [1, -1]])
    return _apply_local_transforms(p, (zmat for _ in range(p.ndim)))