def mat_power(
    mat_m,
    p,
    precision=lax.Precision.HIGHEST,
):
    """A simple matrix power method. M^p where p can be TracedValue."""
    power = jnp.eye(mat_m.shape[0], dtype=_MAT_INV_PTH_ROOT_DTYPE)

    def _iter_condition(state):
        i, _, _ = state
        return i > 0

    def _iter_body(state):
        i, power, mat = state

        power = jax.lax.cond(
            i % 2 == 1,
            lambda: jnp.matmul(mat, power, precision=precision),
            lambda: power,
        )
        i //= 2
        mat = jnp.matmul(mat, mat, precision=precision)
        return i, power, mat

    _, result, _ = lax.while_loop(_iter_condition, _iter_body, (p, power, mat_m))
    return result
A simple matrix power method. M^p where p can be TracedValue.
tools/train/scalable_shampoo/distributed_shampoo.py