def power_iteration(
    matrix,
    num_iters=100,
    error_tolerance=1e-6,
    precision=lax.Precision.HIGHEST,
):
    r"""Power iteration algorithm.

    The power iteration algorithm takes a symmetric PSD matrix `A`, and produces
    a scalar `\lambda` , which is the greatest (in absolute value) eigenvalue
    of `A`, and a vector v, which is the corresponding eigenvector of `A`.

    References:
      [Wikipedia, 2021](https://en.wikipedia.org/wiki/Power_iteration)

    Args:
      matrix: the symmetric PSD matrix.
      num_iters: Number of iterations.
      error_tolerance: Iterative exit condition.
      precision: precision XLA related flag, the available options are: a)
        lax.Precision.DEFAULT (better step time, but not precise) b)
        lax.Precision.HIGH (increased precision, slower) c) lax.Precision.HIGHEST
        (best possible precision, slowest)

    Returns:
      eigen vector, eigen value
    """
    matrix_size = matrix.shape[-1]

    def _iter_condition(state):
        i, unused_v, unused_s, unused_s_v, run_step = state
        return jnp.logical_and(i < num_iters, run_step)

    def _iter_body(state):
        """One step of power iteration."""
        i, new_v, s, s_v, unused_run_step = state
        new_v = new_v / jnp.linalg.norm(new_v)

        s_v = jnp.einsum("ij,j->i", matrix, new_v, precision=precision)
        s_new = jnp.einsum("i,i->", new_v, s_v, precision=precision)
        return (
            i + 1,
            s_v,
            s_new,
            s_v,
            jnp.greater(jnp.abs(s_new - s), error_tolerance),
        )

    # Figure out how to use step as seed for random.
    v_0 = (
        np.random.RandomState(1729).uniform(-1.0, 1.0, matrix_size).astype(matrix.dtype)
    )

    init_state = tuple([0, v_0, jnp.zeros([], dtype=matrix.dtype), v_0, True])
    _, v_out, s_out, _, _ = lax.while_loop(_iter_condition, _iter_body, init_state)
    v_out = v_out / jnp.linalg.norm(v_out)
    return v_out, s_out
Power iteration algorithm.

    The power iteration algorithm takes a symmetric PSD matrix `A`, and produces
    a scalar `\lambda` , which is the greatest (in absolute value) eigenvalue
    of `A`, and a vector v, which is the corresponding eigenvector of `A`.

    References:
      [Wikipedia, 2021](https://en.wikipedia.org/wiki/Power_iteration)

    Args:
      matrix: the symmetric PSD matrix.
      num_iters: Number of iterations.
      error_tolerance: Iterative exit condition.
      precision: precision XLA related flag, the available options are: a)
        lax.Precision.DEFAULT (better step time, but not precise) b)
        lax.Precision.HIGH (increased precision, slower) c) lax.Precision.HIGHEST
        (best possible precision, slowest)

    Returns:
      eigen vector, eigen value
    
tools/train/scalable_shampoo/distributed_shampoo.py