code: def _power_iteration(weight, u, v, eps):
      u = weight @ v
      u_norm = jnp.sqrt(jnp.sum(u**2))
      u = u / jnp.maximum(eps, u_norm)
  
      v = weight.T @ u
      v_norm = jnp.sqrt(jnp.sum(v**2))
      v = v / jnp.maximum(eps, v_norm)
  
      return u, v
  
file_path: equinox/equinox/nn/_spectral_norm.py
imports: from typing import Generic, Optional, TypeVar
  import jax
  import jax.lax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, Float, PRNGKeyArray
  from .._module import field
  from .._tree import tree_at
  from ._sequential import StatefulLayer
  from ._stateful import State, StateIndex
label_desc: Use image node_id 1 for calling a global function with example usage: _power_iteration(weight, u, v, eps) and returns: u,  v