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