optax 0.2.0__tar.gz → 0.2.2__tar.gz

{optax-0.2.0 → optax-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: optax
-Version: 0.2.0
+Version: 0.2.2
 Summary: A gradient processing and optimisation library in JAX.
 Keywords: python,machine learning,reinforcement-learning
 Author-email: Google DeepMind <optax-dev@google.com>
@@ -17,7 +17,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Dist: absl-py>=0.7.1
-Requires-Dist: chex>=0.1.7
+Requires-Dist: chex>=0.1.86
 Requires-Dist: jax>=0.1.55
 Requires-Dist: jaxlib>=0.1.37
 Requires-Dist: numpy>=1.18.0
@@ -33,6 +33,7 @@ Requires-Dist: sphinx-collections>=0.0.1 ; extra == "docs"
 Requires-Dist: tensorflow>=2.4.0 ; extra == "docs"
 Requires-Dist: tensorflow-datasets>=4.2.0 ; extra == "docs"
 Requires-Dist: flax ; extra == "docs"
+Requires-Dist: sphinx_contributors ; extra == "docs"
 Requires-Dist: absl-py>=1.0.0 ; extra == "dp-accounting"
 Requires-Dist: attrs>=21.4.0 ; extra == "dp-accounting"
 Requires-Dist: mpmath>=1.2.1 ; extra == "dp-accounting"

{optax-0.2.0 → optax-0.2.2}/optax/__init__.py

@@ -40,6 +40,7 @@ from optax._src.alias import nadamw
 from optax._src.alias import noisy_sgd
 from optax._src.alias import novograd
 from optax._src.alias import optimistic_gradient_descent
+from optax._src.alias import polyak_sgd
 from optax._src.alias import radam
 from optax._src.alias import rmsprop
 from optax._src.alias import rprop
@@ -115,6 +116,7 @@ from optax._src.transform import scale_by_novograd
 from optax._src.transform import scale_by_optimistic_gradient
 from optax._src.transform import scale_by_param_block_norm
 from optax._src.transform import scale_by_param_block_rms
+from optax._src.transform import scale_by_polyak
 from optax._src.transform import scale_by_radam
 from optax._src.transform import scale_by_rms
 from optax._src.transform import scale_by_rprop
@@ -213,7 +215,7 @@ DifferentiallyPrivateAggregateState = (
 )
 dpsgd = contrib.dpsgd
 
-__version__ = "0.2.0"
+__version__ = "0.2.2"
 
 __all__ = (
     "adabelief",
@@ -308,6 +310,7 @@ __all__ = (
     "piecewise_interpolate_schedule",
     "polynomial_schedule",
     "power_iteration",
+    "polyak_sgd",
     "radam",
     "rmsprop",
     "rprop",
@@ -326,6 +329,7 @@ __all__ = (
     "scale_by_novograd",
     "scale_by_param_block_norm",
     "scale_by_param_block_rms",
+    "scale_by_polyak",
     "scale_by_radam",
     "scale_by_rms",
     "scale_by_rprop",

{optax-0.2.0 → optax-0.2.2}/optax/_src/alias.py

@@ -36,15 +36,43 @@ def adabelief(
     b2: float = 0.999,
     eps: float = 1e-16,
     eps_root: float = 1e-16) -> base.GradientTransformation:
-  """The AdaBelief optimizer.
+  r"""The AdaBelief optimizer.
 
   AdaBelief is an adaptive learning rate optimizer that focuses on fast
   convergence, generalization, and stability. It adapts the step size depending
   on its "belief" in the gradient direction — the optimizer adaptively scales
   the step size by the difference between the predicted and observed gradients.
-  AdaBelief is a modified version of Adam and contains the same number of
-  parameters.
-  
+  AdaBelief is a modified version of :func:`optax.adam` and contains the same
+  number of parameters.
+
+  Let :math:`\alpha_t` represent the learning rate and :math:`\beta_1, \beta_2`,
+  :math:`\varepsilon`, :math:`\bar{\varepsilon}` represent the arguments
+  ``b1``, ``b2``, ``eps`` and ``eps_root`` respectively. The learning rate is
+  indexed by :math:`t` since the learning rate may also be provided by a
+  schedule function.
+
+  The ``init`` function of this optimizer initializes an internal state
+  :math:`S_0 := (m_0, s_0) = (0, 0)`, representing initial estimates for the
+  first and second moments. In practice these values are stored as pytrees
+  containing all zeros, with the same shape as the model updates.
+  At step :math:`t`, the ``update`` function of this optimizer takes as
+  arguments the incoming gradients :math:`g_t` and optimizer state :math:`S_t`
+  and computes updates :math:`u_t` and new state :math:`S_{t+1}`. Thus, for
+  :math:`t > 0`, we have,
+
+  .. math::
+
+    \begin{align*}
+      m_t &\leftarrow \beta_1 \cdot m_{t-1} + (1-\beta_1) \cdot g_t \\
+      s_t &\leftarrow \beta_2 \cdot s_{t-1} + (1-\beta_2) \cdot (g_t - m_t)^2 
+      + \bar{\varepsilon} \\
+      \hat{m}_t &\leftarrow m_t / {(1-\beta_1^t)} \\
+      \hat{s}_t &\leftarrow s_t / {(1-\beta_2^t)} \\
+      u_t &\leftarrow -\alpha_t \cdot \hat{m}_t / \left(\sqrt{\hat{s}_t}
+      + \varepsilon \right) \\
+      S_t &\leftarrow (m_t, s_t).
+    \end{align*}
+
   Examples:
     >>> import optax
     >>> import jax
@@ -322,7 +350,7 @@ def adam(
 
   Let :math:`\alpha_t` represent the learning rate and :math:`\beta_1, \beta_2`,
   :math:`\varepsilon`, :math:`\bar{\varepsilon}` represent the arguments
-  ``b1``, ``b2``, ``eps`` and ``eps_root`` respectievly. The learning rate is
+  ``b1``, ``b2``, ``eps`` and ``eps_root`` respectively. The learning rate is
   indexed by :math:`t` since the learning rate may also be provided by a
   schedule function.
 
@@ -994,9 +1022,27 @@ def noisy_sgd(
 ) -> base.GradientTransformation:
   r"""A variant of SGD with added noise.
 
-  It has been found that adding noise to the gradients can improve
-  both the training error and the generalization error in very deep networks.
-  
+  Noisy SGD is a variant of :func:`optax.sgd` that incorporates Gaussian noise
+  into the updates. It has been found that adding noise to the gradients can
+  improve both the training error and the generalization error in very deep
+  networks.
+
+  The update :math:`u_t` is modified to include this noise as follows:
+
+  .. math::
+    u_t \leftarrow -\alpha_t (g_t + N(0, \sigma_t^2)),
+
+  where :math:`N(0, \sigma_t^2)` represents Gaussian noise with zero mean and a
+  variance of :math:`\sigma_t^2`.
+
+  The variance of this noise decays over time according to the formula
+
+  .. math::
+    \sigma_t^2 = \frac{\eta}{(1+t)^\gamma},
+
+  where :math:`\gamma` is the decay rate parameter ``gamma`` and :math:`\eta`
+  represents the initial variance ``eta``.
+
   Examples:
     >>> import optax
     >>> import jax
@@ -1025,8 +1071,8 @@ def noisy_sgd(
     learning_rate: A global scaling factor, either fixed or evolving along
       iterations with a scheduler, see :func:`optax.scale_by_learning_rate`.
     eta: Initial variance for the Gaussian noise added to gradients.
-    gamma: A parameter controlling the annealing of noise over time, the
-      variance decays according to `(1+t)^-\gamma`.
+    gamma: A parameter controlling the annealing of noise over time ``t``, the
+      variance decays according to ``(1+t)**(-gamma)``.
     seed: Seed for the pseudo-random generation process.
 
   Returns:
@@ -1631,7 +1677,7 @@ def rprop(
   gradient descent. It responds only to the sign of the gradient by increasing
   or decreasing the step size selected per parameter exponentially to speed up
   convergence and avoid oscillations.
-  
+
   Examples:
     >>> import optax
     >>> import jax
@@ -1684,3 +1730,82 @@ def rprop(
       ),
       transform.scale(-1.0),
   )
+
+
+def polyak_sgd(
+    max_learning_rate: float = 1.,
+    scaling: base.ScalarOrSchedule = 1.,
+    f_min: float = 0.0,
+    eps: float = 0.0,
+) -> base.GradientTransformationExtraArgs:
+  r"""SGD with Polyak step-size.
+
+  This solver implements the SGD with Polyak step size of (Loizou et al. 2021).
+  It sets the step-size as
+
+  .. math::
+    s \min\left\{\frac{f(x) - f^\star}{\|\nabla f(x)\|^2 + \epsilon},
+      \gamma_{\max}\right\}\,,
+
+  where :math:`f` is the function from which a gradient is computed,
+  :math:`\gamma_{\max}` is a maximal acceptable learning rate set  by
+  ``max_learning_rate``, :math:`\epsilon` is a constant preventing division by
+  zero set with ``eps``, :math:`s` scales the formula by ``scaling``, and
+  :math:`f^\star` is a guess of the minimum value of the function set with
+  ``f_min``.
+
+
+  Examples:
+    >>> import optax
+    >>> import jax
+    >>> import jax.numpy as jnp
+    >>> def f(x): return jnp.sum(x ** 2)  # simple quadratic function
+    >>> solver = optax.polyak_sgd()
+    >>> params = jnp.array([1., 2., 3.])
+    >>> print('Objective function: ', f(params))
+    Objective function:  14.0
+    >>> opt_state = solver.init(params)
+    >>> for _ in range(5):
+    ...  value, grad = jax.value_and_grad(f)(params)
+    ...  params, opt_state = solver.update(grad, opt_state, params, value=value)
+    ...  print('Objective function: ', f(params))
+    Objective function:  3.5
+    Objective function:  0.875
+    Objective function:  0.21875
+    Objective function:  0.0546875
+    Objective function:  0.013671875
+
+  .. warning::
+    This method requires knowledge of an approximate value of the of the
+    objective function minimum, passed through the ``f_min`` argument.
+    For models that interpolate the data, this can be set to 0 (default
+    value).
+    Failing to set an appropriate value for ``f_min`` can lead to
+    divergence or convergence to a suboptimal solution.
+
+  References:
+    Loizou et al. `Stochastic polyak step-size for SGD: An adaptive learning
+    rate for fast convergence <https://arxiv.org/abs/2002.10542>`_, 2021
+
+    Berrada et al., `Training neural networks for and by interpolation
+    <https://arxiv.org/pdf/1906.05661.pdf>`_, 2020
+
+  Args:
+    max_learning_rate: a maximum step size to use (defaults to 1).
+    scaling: A global scaling factor, either fixed or evolving along
+      iterations with a scheduler (defaults to 1).
+    f_min: a lower bound on the objective function (defaults to 0). Corresponds
+      to :math:`f^\star` in the formula above.
+    eps: a value to add in the denominator of the update (defaults to 0).
+
+  Returns:
+    A :class:`GradientTransformationExtraArgs`, where the ``update`` function
+    takes an additional keyword argument ``value`` containing the current
+    value of the objective function.
+  """
+  return combine.chain(
+      sgd(learning_rate=scaling),
+      transform.scale_by_polyak(
+          max_learning_rate=max_learning_rate, f_min=f_min, eps=eps
+      ),
+  )

{optax-0.2.0 → optax-0.2.2}/optax/_src/alias_test.py

@@ -60,6 +60,7 @@ _OPTIMIZERS_UNDER_TEST = (
     dict(opt_name='rprop', opt_kwargs=dict(learning_rate=1e-1)),
     dict(opt_name='sm3', opt_kwargs=dict(learning_rate=1.0)),
     dict(opt_name='yogi', opt_kwargs=dict(learning_rate=1e-1)),
+    dict(opt_name='polyak_sgd', opt_kwargs=dict(max_learning_rate=1.))
 )
 
 
@@ -71,11 +72,10 @@ def _setup_parabola(dtype):
   if jnp.iscomplexobj(dtype):
     final_params *= 1 + 1j
 
-  @jax.grad
-  def get_updates(params):
+  def objective(params):
     return jnp.sum(numerics.abs_sq(params - final_params))
 
-  return initial_params, final_params, get_updates
+  return initial_params, final_params, objective
 
 
 def _setup_rosenbrock(dtype):
@@ -89,12 +89,11 @@ def _setup_rosenbrock(dtype):
   initial_params = jnp.array([0.0, 0.0], dtype=dtype)
   final_params = jnp.array([a, a**2], dtype=dtype)
 
-  @jax.grad
-  def get_updates(params):
+  def objective(params):
     return (numerics.abs_sq(a - params[0]) +
             b * numerics.abs_sq(params[1] - params[0]**2))
 
-  return initial_params, final_params, get_updates
+  return initial_params, final_params, objective
 
 
 class AliasTest(chex.TestCase):
@@ -113,21 +112,26 @@ class AliasTest(chex.TestCase):
         'lion',
         'rprop',
         'adadelta',
+        'polyak_sgd',
     ) and jnp.iscomplexobj(dtype):
       raise absltest.SkipTest(
           f'{opt_name} does not support complex parameters.'
       )
 
     opt = getattr(alias, opt_name)(**opt_kwargs)
-    initial_params, final_params, get_updates = target(dtype)
+    initial_params, final_params, objective = target(dtype)
 
     @jax.jit
     def step(params, state):
-      updates = get_updates(params)
+      value, updates = jax.value_and_grad(objective)(params)
       # Complex gradients need to be conjugated before being added to parameters
       # https://gist.github.com/wdphy16/118aef6fb5f82c49790d7678cf87da29
       updates = jax.tree_util.tree_map(lambda x: x.conj(), updates)
-      updates, state = opt.update(updates, state, params)
+      if opt_name == 'polyak_sgd':
+        update_kwargs = {'value': value}
+      else:
+        update_kwargs = {}
+      updates, state = opt.update(updates, state, params, **update_kwargs)
       params = update.apply_updates(params, updates)
       return params, state
 
@@ -146,13 +150,13 @@ class AliasTest(chex.TestCase):
   def test_optimizers_can_be_wrapped_in_inject_hyperparams(
       self, opt_name, opt_kwargs):
     """Checks that optimizers can be wrapped in inject_hyperparams."""
-    # See also https://github.com/deepmind/optax/issues/412.
+    # See also https://github.com/google-deepmind/optax/issues/412.
     opt_factory = getattr(alias, opt_name)
     opt = opt_factory(**opt_kwargs)
     if opt_name == 'adafactor':
       # Adafactor wrapped in inject_hyperparams currently needs a static
       # argument to be specified in order to be jittable. See issue
-      # https://github.com/deepmind/optax/issues/412.
+      # https://github.com/google-deepmind/optax/issues/412.
       opt_inject = _inject.inject_hyperparams(
           opt_factory, static_args=('min_dim_size_to_factor',))(**opt_kwargs)
     else:
@@ -162,11 +166,17 @@ class AliasTest(chex.TestCase):
     grads = [jnp.ones((2, 3)), jnp.negative(jnp.ones((2, 5, 2)))]
 
     state = self.variant(opt.init)(params)
-    updates, new_state = self.variant(opt.update)(grads, state, params)
+    if opt_name == 'polyak_sgd':
+      update_kwargs = {'value': jnp.array(0.)}
+    else:
+      update_kwargs = {}
+    updates, new_state = self.variant(opt.update)(
+        grads, state, params, **update_kwargs
+    )
 
     state_inject = self.variant(opt_inject.init)(params)
     updates_inject, new_state_inject = self.variant(opt_inject.update)(
-        grads, state_inject, params)
+        grads, state_inject, params, **update_kwargs)
 
     with self.subTest('Equality of updates.'):
       chex.assert_trees_all_close(updates_inject, updates, rtol=1e-4)

{optax-0.2.0 → optax-0.2.2}/optax/_src/combine.py

@@ -27,16 +27,30 @@ def chain(
 ) -> base.GradientTransformationExtraArgs:
   """Applies a list of chainable update transformations.
 
-  Given a sequence of chainable transforms, `chain` returns an `init_fn`
-  that constructs a `state` by concatenating the states of the individual
-  transforms, and returns an `update_fn` which chains the update transformations
-  feeding the appropriate state to each.
+  This function creates a new :func:`optax.GradientTransformation` that applies
+  a sequence of gradient transformations in order. The ``init`` function of the
+  new transformation constructs the optimizer state by concatenating the states
+  of the individual transforms, while the ``update`` function applies the
+  updates in the given order.
+
+  Examples:
+
+    A transform that scales by -0.1 the adam update:
+
+      >>> import optax
+      >>> transform1 = optax.scale_by_adam()
+      >>> transform2 = optax.scale(-0.1)
+      >>> chained_transform = optax.chain(transform1, transform2)
+      >>> params = {'a': 1.0}
+      >>> state = chained_transform.init(params)
+      >>> updates = {'a': -0.5}
+      >>> updates, new_state = chained_transform.update(updates, state, params)
 
   Args:
     *args: a sequence of chainable (init_fn, update_fn) tuples.
 
   Returns:
-    A ``GradientTransformationExtraArgs``, created by chaining the input
+    A :func:`GradientTransformationExtraArgs`, created by chaining the input
     transformations. Note that independent of the argument types, the resulting
     transformation always supports extra args. Any extra arguments passed to the
     returned transformation will be passed only to those transformations in the
@@ -147,7 +161,7 @@ def multi_transform(
     import jax.numpy as jnp
 
     def map_nested_fn(fn):
-      '''Recursively apply `fn` to the key-value pairs of a nested dict'''
+      '''Recursively apply `fn` to the key-value pairs of a nested dict.'''
       def map_fn(nested_dict):
         return {k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
                 for k, v in nested_dict.items()}
@@ -178,7 +192,7 @@ def multi_transform(
         param_labels)
 
   If you would like to not optimize some parameters, you may wrap
-  ``optax.multi_transform`` with :func:`optax.masked`.
+  :func:`optax.multi_transform` with :func:`optax.masked`.
 
   Args:
     transforms: A mapping from labels to transformations. Each transformation
@@ -191,7 +205,8 @@ def multi_transform(
       extra_arg fields with the same tree structure as params/updates.
 
   Returns:
-    An ``optax.GradientTransformation``.
+    A :func:`optax.GradientTransformationExtraArgs` that implements an ``init``
+    and ``update`` function.
   """
 
   transforms = {

{optax-0.2.0 → optax-0.2.2}/optax/_src/linear_algebra.py

@@ -14,70 +14,126 @@
 # ==============================================================================
 """Linear algebra utilities used in optimisation."""
 
+import functools
+from typing import Callable, Optional, Union
+
 import chex
 import jax
 from jax import lax
 import jax.numpy as jnp
-import numpy as np
-
+from optax import tree_utils as otu
 from optax._src import base
 from optax._src import numerics
 
 
+def _normalize_tree(x):
+  # divide by the L2 norm of the tree weights.
+  return otu.tree_scalar_mul(1.0 / otu.tree_l2_norm(x), x)
+
+
 def global_norm(updates: base.PyTree) -> chex.Array:
   """Compute the global norm across a nested structure of tensors."""
   return jnp.sqrt(sum(
       jnp.sum(numerics.abs_sq(x)) for x in jax.tree_util.tree_leaves(updates)))
 
 
-def power_iteration(matrix: chex.Array,
-                    num_iters: int = 100,
-                    error_tolerance: float = 1e-6,
-                    precision: lax.Precision = lax.Precision.HIGHEST):
+def _power_iteration_cond_fun(error_tolerance, num_iters, loop_vars):
+  normalized_eigvec, unnormalized_eigvec, eig, iter_num = loop_vars
+  residual = otu.tree_sub(
+      unnormalized_eigvec, otu.tree_scalar_mul(eig, normalized_eigvec)
+  )
+  residual_norm = otu.tree_l2_norm(residual)
+  converged = jnp.abs(residual_norm / eig) < error_tolerance
+  return ~converged & (iter_num < num_iters)
+
+
+def power_iteration(
+    matrix: Union[chex.Array, Callable[[chex.ArrayTree], chex.ArrayTree]],
+    *,
+    v0: Optional[chex.ArrayTree] = None,
+    num_iters: int = 100,
+    error_tolerance: float = 1e-6,
+    precision: lax.Precision = lax.Precision.HIGHEST,
+    key: Optional[chex.PRNGKey] = None,
+) -> tuple[chex.Numeric, chex.ArrayTree]:
   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`.
+  This algorithm computes the dominant eigenvalue and its associated eigenvector
+  of a diagonalizable matrix. This matrix can be given as an array or as a
+  callable implementing a matrix-vector product.
 
   References:
-    [Wikipedia, 2021](https://en.wikipedia.org/wiki/Power_iteration)
+    Wikipedia contributors. `Power iteration
+    <https://en.wikipedia.org/w/index.php?tit0le=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).
+    matrix: a square matrix, either as an array or a callable implementing a
+      matrix-vector product.
+    v0: initial vector approximating the dominiant eigenvector. If ``matrix``
+      is an array of size (n, n), v0 must be a vector of size (n,). If instead
+      ``matrix`` is a callable, then v0 must be a tree with the same structure
+      as the input of this callable. If this argument is None and ``matrix`` is
+      an array, then a random vector sampled from a uniform distribution in
+      [-1, 1] is used as initial vector.
+    num_iters: Number of power iterations.
+    error_tolerance: Iterative exit condition. The procedure stops when the
+      relative error of the estimate of the dominant eigenvalue is below this
+      threshold.
+    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).
+    key: random key for the initialization of ``v0`` when not given
+      explicitly. When this argument is None, `jax.random.PRNGKey(0)` is used.
 
   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)
+    A pair (eigenvalue, eigenvector), where eigenvalue is the dominant
+    eigenvalue of ``matrix`` and eigenvector is its associated eigenvector.
 
-    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.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
+  .. versionchanged:: 0.2.2
+    ``matrix`` can be a callable. Reversed the order of the return parameters,
+    from (eigenvector, eigenvalue) to (eigenvalue, eigenvector).
+  """
+  if callable(matrix):
+    mvp = matrix
+    if v0 is None:
+      # v0 must be given as we don't know the underlying pytree structure.
+      raise ValueError('v0 must be provided when `matrix` is a callable.')
+  else:
+    mvp = lambda v: jnp.matmul(matrix, v, precision=precision)
+    if v0 is None:
+      if key is None:
+        key = jax.random.PRNGKey(0)
+      # v0 is uniformly distributed in [-1, 1]
+      v0 = jax.random.uniform(
+          key,
+          shape=matrix.shape[-1:],
+          dtype=matrix.dtype,
+          minval=-1.0,
+          maxval=1.0,
+      )
+
+  v0 = _normalize_tree(v0)
+
+  cond_fun = functools.partial(
+      _power_iteration_cond_fun,
+      error_tolerance,
+      num_iters,
+  )
+
+  def _body_fun(loop_vars):
+    _, z, _, iter_num = loop_vars
+    eigvec = _normalize_tree(z)
+    z = mvp(eigvec)
+    eig = otu.tree_vdot(eigvec, z)
+    return eigvec, z, eig, iter_num + 1
+
+  init_vars = (v0, mvp(v0), jnp.asarray(0.0), jnp.asarray(0))
+  _, unormalized_eigenvector, dominant_eigenvalue, _ = (
+      jax.lax.while_loop(cond_fun, _body_fun, init_vars)
+  )
+  normalized_eigenvector = _normalize_tree(unormalized_eigenvector)
+  return dominant_eigenvalue, normalized_eigenvector
 
 
 def matrix_inverse_pth_root(matrix: chex.Array,
@@ -117,7 +173,7 @@ def matrix_inverse_pth_root(matrix: chex.Array,
   matrix_size = matrix.shape[0]
   alpha = jnp.asarray(-1.0 / p, jnp.float32)
   identity = jnp.eye(matrix_size, dtype=jnp.float32)
-  _, max_ev = power_iteration(
+  max_ev, _ = power_iteration(
       matrix=matrix, num_iters=100,
       error_tolerance=1e-6, precision=precision)
   ridge_epsilon = ridge_epsilon * jnp.maximum(max_ev, 1e-16)