# Path: optax/_src/clipping.py
# def clip(max_delta: chex.Numeric) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def clip_by_block_rms(threshold: float) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
#     def _clip_fn(u):
# def clip_by_global_norm(max_norm: float) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def unitwise_norm(x: chex.Array) -> chex.Array:
# def unitwise_clip(g_norm: chex.Array,
#                   max_norm: chex.Array,
#                   grad: chex.Array,
#                   div_eps: float = 1e-6) -> chex.Array:
# def adaptive_grad_clip(clipping: float,
#                        eps: float = 1e-3) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params):
# 
# Path: optax/_src/linear_algebra.py
# def global_norm(updates: base.Updates) -> base.Updates:
# def power_iteration(
#     matrix,
#     num_iters=100,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _iter_condition(state):
#   def _iter_body(state):
# def matrix_inverse_pth_root(
#     matrix,
#     p,
#     num_iters=100,
#     ridge_epsilon=1e-6,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _unrolled_mat_pow_1(mat_m):
#   def _unrolled_mat_pow_2(mat_m):
#   def _unrolled_mat_pow_4(mat_m):
#   def _unrolled_mat_pow_8(mat_m):
#   def mat_power(mat_m, p):
#   def _iter_condition(state):
#   def _iter_body(state):


from absl.testing import absltest
from optax._src import clipping
from optax._src import linear_algebra
import chex
import jax
import jax.numpy as jnp
      clipper = clipping.clip_by_block_rms(1. / i)
      # Check that the clipper actually works and block rms is <= threshold
      updates, _ = clipper.update(updates, None)
      self.assertAlmostEqual(rmf_fn(updates[0]), 1. / i)
      self.assertAlmostEqual(rmf_fn(updates[1]), 1. / i)
      # Check that continuously clipping won't cause numerical issues.
      updates_step, _ = clipper.update(self.per_step_updates, None)
      chex.assert_tree_all_close(updates, updates_step)

  def test_clip_by_global_norm(self):
    updates = self.per_step_updates
    for i in range(1, STEPS + 1):
      clipper = clipping.clip_by_global_norm(1. / i)
      # Check that the clipper actually works and global norm is <= max_norm
      updates, _ = clipper.update(updates, None)
      self.assertAlmostEqual(
          linear_algebra.global_norm(updates), 1. / i, places=6)
      # Check that continuously clipping won't cause numerical issues.
      updates_step, _ = clipper.update(self.per_step_updates, None)
      chex.assert_tree_all_close(updates, updates_step)

  def test_adaptive_grad_clip(self):
    updates = self.per_step_updates
    params = self.init_params
    for i in range(1, STEPS + 1):
      clip_r = 1. / i
      clipper = clipping.adaptive_grad_clip(clip_r)

      # Check that the clipper actually works and upd_norm is < c * param_norm.
      updates, _ = clipper.update(updates, None, params)

# Path: optax/_src/clipping.py
# def clip(max_delta: chex.Numeric) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def clip_by_block_rms(threshold: float) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
#     def _clip_fn(u):
# def clip_by_global_norm(max_norm: float) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def unitwise_norm(x: chex.Array) -> chex.Array:
# def unitwise_clip(g_norm: chex.Array,
#                   max_norm: chex.Array,
#                   grad: chex.Array,
#                   div_eps: float = 1e-6) -> chex.Array:
# def adaptive_grad_clip(clipping: float,
#                        eps: float = 1e-3) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params):
# 
# Path: optax/_src/linear_algebra.py
# def global_norm(updates: base.Updates) -> base.Updates:
# def power_iteration(
#     matrix,
#     num_iters=100,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _iter_condition(state):
#   def _iter_body(state):
# def matrix_inverse_pth_root(
#     matrix,
#     p,
#     num_iters=100,
#     ridge_epsilon=1e-6,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _unrolled_mat_pow_1(mat_m):
#   def _unrolled_mat_pow_2(mat_m):
#   def _unrolled_mat_pow_4(mat_m):
#   def _unrolled_mat_pow_8(mat_m):
#   def mat_power(mat_m, p):
#   def _iter_condition(state):
#   def _iter_body(state):


# Path: optax/_src/clipping_test.py
from absl.testing import absltest
from optax._src import clipping
from optax._src import linear_algebra
import chex
import jax
import jax.numpy as jnp

      clipper = clipping.clip_by_block_rms(1. / i)
      # Check that the clipper actually works and block rms is <= threshold
      updates, _ = clipper.update(updates, None)
      self.assertAlmostEqual(rmf_fn(updates[0]), 1. / i)
      self.assertAlmostEqual(rmf_fn(updates[1]), 1. / i)
      # Check that continuously clipping won't cause numerical issues.
      updates_step, _ = clipper.update(self.per_step_updates, None)
      chex.assert_tree_all_close(updates, updates_step)

  def test_clip_by_global_norm(self):
    updates = self.per_step_updates
    for i in range(1, STEPS + 1):
      clipper = clipping.clip_by_global_norm(1. / i)
      # Check that the clipper actually works and global norm is <= max_norm
      updates, _ = clipper.update(updates, None)
      self.assertAlmostEqual(
          linear_algebra.global_norm(updates), 1. / i, places=6)
      # Check that continuously clipping won't cause numerical issues.
      updates_step, _ = clipper.update(self.per_step_updates, None)
      chex.assert_tree_all_close(updates, updates_step)

  def test_adaptive_grad_clip(self):
    updates = self.per_step_updates
    params = self.init_params
    for i in range(1, STEPS + 1):
      clip_r = 1. / i
      clipper = clipping.adaptive_grad_clip(clip_r)

      # Check that the clipper actually works and upd_norm is < c * param_norm.
      updates, _ = clipper.update(updates, None, params)

      u_norm, p_norm = jax.tree_map(clipping.unitwise_norm, (updates, params))