def pow_pauli_combination(
    ai: cirq.TParamValComplex,
    ax: cirq.TParamValComplex,
    ay: cirq.TParamValComplex,
    az: cirq.TParamValComplex,
    exponent: int,
) -> tuple[
    cirq.TParamValComplex, cirq.TParamValComplex, cirq.TParamValComplex, cirq.TParamValComplex
]:
    """Computes non-negative integer power of single-qubit Pauli combination.

    Returns scalar coefficients bi, bx, by, bz such that

        bi I + bx X + by Y + bz Z = (ai I + ax X + ay Y + az Z)^exponent

    Correctness of the formulas below follows from the binomial expansion
    and the fact that for any real or complex vector (ax, ay, az) and any
    non-negative integer k:

         [ax X + ay Y + az Z]^(2k) = (ax^2 + ay^2 + az^2)^k I

    """
    if exponent == 0:
        return 1, 0, 0, 0

    if any(isinstance(a, sympy.Basic) for a in [ax, ay, az]):
        v = sympy.sqrt(ax * ax + ay * ay + az * az)
    else:
        v = np.sqrt(ax * ax + ay * ay + az * az).item()
    s = (ai + v) ** exponent
    t = (ai - v) ** exponent

    ci = (s + t) / 2
    if s == t:
        # v is near zero, only one term in binomial expansion survives
        cxyz = exponent * ai ** (exponent - 1)
    else:
        # v is non-zero, account for all terms of binomial expansion
        cxyz = (s - t) / 2
        cxyz = cxyz / v

    return ci, cxyz * ax, cxyz * ay, cxyz * az
Computes non-negative integer power of single-qubit Pauli combination.

    Returns scalar coefficients bi, bx, by, bz such that

        bi I + bx X + by Y + bz Z = (ai I + ax X + ay Y + az Z)^exponent

    Correctness of the formulas below follows from the binomial expansion
    and the fact that for any real or complex vector (ax, ay, az) and any
    non-negative integer k:

         [ax X + ay Y + az Z]^(2k) = (ax^2 + ay^2 + az^2)^k I

    
cirq-core/cirq/linalg/operator_spaces.py