def map_eigenvalues(
    matrix: np.ndarray, func: Callable[[complex], complex], *, atol: float = 1e-8
) -> np.ndarray:
    """Applies a function to the eigenvalues of a matrix.

    Given M = sum_k a_k |v_k><v_k|, returns f(M) = sum_k f(a_k) |v_k><v_k|.

    Args:
        matrix: The matrix to modify with the function.
        func: The function to apply to the eigenvalues of the matrix.
        atol: Absolute threshold used when separating eigenspaces.

    Returns:
        The transformed matrix.
    """
    vals, vecs = unitary_eig(matrix, atol=atol)
    pieces = [np.outer(vec, np.conj(vec.T)) for vec in vecs.T]
    out_vals = np.vectorize(func)(vals.astype(complex))

    total = np.zeros(shape=matrix.shape)
    for piece, val in zip(pieces, out_vals):
        total = np.add(total, piece * val)
    return total
Applies a function to the eigenvalues of a matrix.

    Given M = sum_k a_k |v_k><v_k|, returns f(M) = sum_k f(a_k) |v_k><v_k|.

    Args:
        matrix: The matrix to modify with the function.
        func: The function to apply to the eigenvalues of the matrix.
        atol: Absolute threshold used when separating eigenspaces.

    Returns:
        The transformed matrix.
    
cirq-core/cirq/linalg/decompositions.py