Decomposes a 2-qubit unitary into 1-qubit ops and XX/YY/ZZ interactions.

    Args:
        mat: The 4x4 unitary matrix to decompose.
        rtol: Per-matrix-entry relative tolerance on equality.
        atol: Per-matrix-entry absolute tolerance on equality.

    Returns:
        A `cirq.KakDecomposition` canonicalized such that the interaction
        coefficients x, y, z satisfy:

            0 ≤ abs(z) ≤ y ≤ x ≤ π/4
            z ≠ -π/4

    Raises:
        ValueError: Bad matrix.
        ArithmeticError: Failed to perform the decomposition.

    References:
        'An Introduction to Cartan's KAK Decomposition for QC Programmers'
        https://arxiv.org/abs/quant-ph/0507171

def kak_decomposition(
        mat: np.ndarray,
        rtol: float = 1e-5,
        atol: float = 1e-8) -> KakDecomposition:
    """Decomposes a 2-qubit unitary into 1-qubit ops and XX/YY/ZZ interactions.

    Args:
        mat: The 4x4 unitary matrix to decompose.
        rtol: Per-matrix-entry relative tolerance on equality.
        atol: Per-matrix-entry absolute tolerance on equality.

    Returns:
        A `cirq.KakDecomposition` canonicalized such that the interaction
        coefficients x, y, z satisfy:

            0 ≤ abs(z) ≤ y ≤ x ≤ π/4
            z ≠ -π/4

    Raises:
        ValueError: Bad matrix.
        ArithmeticError: Failed to perform the decomposition.

    References:
        'An Introduction to Cartan's KAK Decomposition for QC Programmers'
        https://arxiv.org/abs/quant-ph/0507171
    """
    magic = np.array([[1, 0, 0, 1j],
                      [0, 1j, 1, 0],
                      [0, 1j, -1, 0],
                      [1, 0, 0, -1j]]) * np.sqrt(0.5)
    gamma = np.array([[1, 1, 1, 1],
                      [1, 1, -1, -1],
                      [-1, 1, -1, 1],
                      [1, -1, -1, 1]]) * 0.25

    # Diagonalize in magic basis.
    left, d, right = diagonalize.bidiagonalize_unitary_with_special_orthogonals(
        combinators.dot(np.conj(magic.T), mat, magic),
        atol=atol,
        rtol=rtol,
        check_preconditions=False)

    # Recover pieces.
    a1, a0 = so4_to_magic_su2s(left.T,
                               atol=atol,
                               rtol=rtol,
                               check_preconditions=False)
    b1, b0 = so4_to_magic_su2s(right.T,
                               atol=atol,
                               rtol=rtol,
                               check_preconditions=False)
    w, x, y, z = gamma.dot(np.vstack(np.angle(d))).flatten()
    g = np.exp(1j * w)

    # Canonicalize.
    inner_cannon = kak_canonicalize_vector(x, y, z)

    b1 = np.dot(inner_cannon.single_qubit_operations_before[0], b1)
    b0 = np.dot(inner_cannon.single_qubit_operations_before[1], b0)
    a1 = np.dot(a1, inner_cannon.single_qubit_operations_after[0])
    a0 = np.dot(a0, inner_cannon.single_qubit_operations_after[1])
    return KakDecomposition(
        interaction_coefficients=inner_cannon.interaction_coefficients,
        global_phase=g * inner_cannon.global_phase,
        single_qubit_operations_before=(b1, b0),
        single_qubit_operations_after=(a1, a0))