def matrix_inverse_pth_root(
    matrix,
    p,
    num_iters=100,
    ridge_epsilon=1e-6,
    error_tolerance=1e-6,
    precision=lax.Precision.HIGHEST,
    relative_matrix_epsilon=True,
    lobpcg_topk_precondition=0,
    lobpcg_max_iter=0,
):
    """Computes `matrix^(-1/p)`, where `p` is a positive integer.

    This function uses the Coupled newton iterations algorithm for
    the computation of a matrix's inverse pth root.


    References:
      [Functions of Matrices, Theory and Computation,
       Nicholas J Higham, Pg 184, Eq 7.18](
       https://epubs.siam.org/doi/book/10.1137/1.9780898717778)

    Args:
      matrix: the symmetric PSD matrix whose power it to be computed
      p: exponent, for p a positive integer.
      num_iters: Maximum number of iterations.
      ridge_epsilon: Ridge epsilon added to make the matrix positive definite.
      error_tolerance: Error indicator, useful for early termination.
      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)
      relative_matrix_epsilon: Whether to use relative epsilon to the max eigen
        value when computing inverse-pth root.
      lobpcg_topk_precondition: If nonzero, specifies the number of top
        eigenvectors to subtract out before performing LOBPCG. Note this makes
        relative_matrix_epsilon essentially free.
      lobpcg_max_iter: Maximum iteration count for LOBPCG, defaults to
        `lobpcg_topk_precondition`.

    Returns:
      matrix^(-1/p) and the error
    """

    # If the input is not square, materialize it from the concatenated form.
    if matrix.shape[0] != matrix.shape[1]:
        matrix = symmetric_matrices.materialize_matrix_from_concat(matrix)

    assert matrix.shape[0] == matrix.shape[1]

    # We use _MAT_INV_PTH_ROOT_DTYPE for the matrix inverse pth root.
    # Switch to f64 if you have hardware that supports it. Enable the jax flag
    # jax_enable_x64 for this to work.
    matrix_size = matrix.shape[0]
    orig_dtype = matrix.dtype
    matrix = matrix.astype(_MAT_INV_PTH_ROOT_DTYPE)
    alpha = jnp.asarray(-1.0 / p, _MAT_INV_PTH_ROOT_DTYPE)
    identity = jnp.eye(matrix_size, dtype=_MAT_INV_PTH_ROOT_DTYPE)
    original_matrix = matrix

    if lobpcg_topk_precondition > 0:
        # TODO(vladf): reuse previous top-k as the initial search directions
        pad_shape = (matrix_size - lobpcg_topk_precondition, lobpcg_topk_precondition)
        search_dirs = jnp.concatenate(
            (jnp.eye(lobpcg_topk_precondition), jnp.zeros(pad_shape)), axis=0
        )
        eigvals, eigvecs, actual_iters = linalg.lobpcg_standard(
            matrix,
            search_dirs,
            lobpcg_topk_precondition if lobpcg_max_iter == 0 else lobpcg_max_iter,
        )
        del actual_iters  # TODO(vladf): return diagnostics dictionary

        # The minimal eigenvalue among top-k becomes the maximal one in the whole
        # matrix after deflation.
        max_ev = jnp.min(eigvals)
        deflation = eigvals - max_ev
        scaled_vecs = eigvecs * jnp.sqrt(deflation)

        # Deflate out top eigenvectors to reduce matrix condition number.
        matrix -= scaled_vecs.dot(scaled_vecs.T, precision=jax.lax.Precision.HIGHEST)

    # Only use power iteration if lobpcg wasn't already used to derive the
    # top eigenvalue.
    elif relative_matrix_epsilon:
        _, max_ev = power_iteration(
            matrix=matrix, num_iters=100, error_tolerance=1e-6, precision=precision
        )
        eigvals, eigvecs = None, None  # Unused but required by pytype.

    # Use absolute matrix epsilon scaling otherwise.
    else:
        max_ev = 1.0
        eigvals, eigvecs = None, None  # Unused but required by pytype.

    ridge_epsilon = ridge_epsilon * jnp.maximum(max_ev, 1e-6)

    def _iter_condition(state):
        (i, unused_mat_m, unused_mat_h, unused_old_mat_h, error, run_step) = state
        error_above_threshold = jnp.logical_and(error > error_tolerance, run_step)
        return jnp.logical_and(i < num_iters, error_above_threshold)

    def _iter_body(state):
        (i, mat_m, mat_h, unused_old_mat_h, error, unused_run_step) = state
        mat_m_i = (1 - alpha) * identity + alpha * mat_m
        new_mat_m = jnp.matmul(mat_power(mat_m_i, p), mat_m, precision=precision)
        new_mat_h = jnp.matmul(mat_h, mat_m_i, precision=precision)
        new_error = jnp.max(jnp.abs(new_mat_m - identity))
        # sometimes error increases after an iteration before decreasing and
        # converging. 1.2 factor is used to bound the maximal allowed increase.
        return (i + 1, new_mat_m, new_mat_h, mat_h, new_error, new_error < error * 1.2)

    if matrix_size == 1:
        resultant_mat_h = (matrix + ridge_epsilon) ** alpha
        error = jnp.array(0, jnp.float32)
    else:
        damped_matrix = matrix + ridge_epsilon * identity

        z = (1 + p) / (2 * jnp.linalg.norm(damped_matrix))
        new_mat_m_0 = damped_matrix * z
        new_error = jnp.max(jnp.abs(new_mat_m_0 - identity))
        new_mat_h_0 = identity * jnp.power(z, 1.0 / p)
        init_state = tuple([0, new_mat_m_0, new_mat_h_0, new_mat_h_0, new_error, True])
        _, mat_m, mat_h, old_mat_h, error, convergence = lax.while_loop(
            _iter_condition, _iter_body, init_state
        )
        error = jnp.max(jnp.abs(mat_m - identity)).astype(jnp.float32)
        is_converged = jnp.asarray(convergence, old_mat_h.dtype)
        resultant_mat_h = is_converged * mat_h + (1 - is_converged) * old_mat_h
        resultant_mat_h = jnp.asarray(resultant_mat_h, orig_dtype)

    if lobpcg_topk_precondition > 0:
        # Since we deflated the top eigenvectors prior to p-th root inverse,
        # the resultant matrix has larger eigenvalues associated with those
        # same eigenvectors, which we need to now re-deflate.
        #
        # Note that _pth_root_difference returns positive values for this
        # particular argument ordering as min(eigvals) <= eigvals for the
        # jnp.sqrt below.
        pth_diff = _pth_root_difference(ridge_epsilon, jnp.min(eigvals), eigvals, p)
        scaled_vecs = eigvecs * jnp.sqrt(pth_diff)
        resultant_mat_h = (
            resultant_mat_h.astype(scaled_vecs.dtype)
            - scaled_vecs.dot(scaled_vecs.T, precision=jax.lax.Precision.HIGHEST)
        ).astype(orig_dtype)
        mat_m = jnp.matmul(
            mat_power(resultant_mat_h, p),
            original_matrix,
            precision=jax.lax.Precision.HIGHEST,
        )
        error = jnp.max(jnp.abs(mat_m - identity)).astype(jnp.float32)

    return resultant_mat_h, error
Computes `matrix^(-1/p)`, where `p` is a positive integer.

    This function uses the Coupled newton iterations algorithm for
    the computation of a matrix's inverse pth root.


    References:
      [Functions of Matrices, Theory and Computation,
       Nicholas J Higham, Pg 184, Eq 7.18](
       https://epubs.siam.org/doi/book/10.1137/1.9780898717778)

    Args:
      matrix: the symmetric PSD matrix whose power it to be computed
      p: exponent, for p a positive integer.
      num_iters: Maximum number of iterations.
      ridge_epsilon: Ridge epsilon added to make the matrix positive definite.
      error_tolerance: Error indicator, useful for early termination.
      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)
      relative_matrix_epsilon: Whether to use relative epsilon to the max eigen
        value when computing inverse-pth root.
      lobpcg_topk_precondition: If nonzero, specifies the number of top
        eigenvectors to subtract out before performing LOBPCG. Note this makes
        relative_matrix_epsilon essentially free.
      lobpcg_max_iter: Maximum iteration count for LOBPCG, defaults to
        `lobpcg_topk_precondition`.

    Returns:
      matrix^(-1/p) and the error
    
tools/train/scalable_shampoo/distributed_shampoo.py