# Path: optax/_src/base.py
# NO_PARAMS_MSG = (
#     'You are using a transformation that requires the current value of '
#     'parameters, but you are not passing `params` when calling `update`.')
# class TransformInitFn(typing_extensions.Protocol):
# class TransformUpdateFn(typing_extensions.Protocol):
# class GradientTransformation(NamedTuple):
# class EmptyState(NamedTuple):
#   def __call__(self, params: Params) -> OptState:
#   def __call__(
#       self,
#       updates: Updates,
#       state: OptState,
#       params: Optional[Params] = None
#     ) -> Tuple[Updates, OptState]:
# def identity() -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def set_to_zero() -> GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def stateless(
#     f: Callable[[Updates, Optional[Params]], Updates],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def stateless_with_tree_map(
#     f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# 
# 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/numerics.py
# def abs_sq(x: chex.Array) -> chex.Array:
# def safe_norm(x: chex.Array,
#               min_norm: chex.Numeric,
#               ord: Optional[Union[int, float, str]] = None,  # pylint: disable=redefined-builtin
#               axis: Union[None, Tuple[int, ...], int] = None,
#               keepdims: bool = False) -> chex.Array:
# def safe_root_mean_squares(x: chex.Array, min_rms: chex.Numeric) -> chex.Array:
# def safe_int32_increment(count: chex.Numeric) -> chex.Numeric:


import chex
import jax
import jax.numpy as jnp
from optax._src import base
from optax._src import linear_algebra
from optax._src import numerics
    del params
    return ClipByGlobalNormState()

  def update_fn(updates, state, params=None):
    del params
    g_norm = linear_algebra.global_norm(updates)
    # TODO(b/163995078): revert back to the following (faster) implementation
    # once analysed how it affects backprop through update (e.g. meta-gradients)
    # g_norm = jnp.maximum(max_norm, g_norm)
    # updates = jax.tree_map(lambda t: (t / g_norm) * max_norm, updates)
    trigger = jnp.squeeze(g_norm < max_norm)
    chex.assert_shape(trigger, ())  # A scalar.

    updates = jax.tree_map(
        lambda t: jax.lax.select(trigger, t, (t / g_norm) * max_norm), updates)
    return updates, state

  return base.GradientTransformation(init_fn, update_fn)


def unitwise_norm(x: chex.Array) -> chex.Array:
  """Computes norms of each output unit separately."""
  if jnp.squeeze(x).ndim <= 1:  # Scalars and vectors
    squared_norm = jnp.sum(numerics.abs_sq(x), keepdims=True)
  # Note that this assumes parameters with a shape of length 3 are multihead
  # linear parameters--if you wish to apply AGC to 1D convs, you may need
  # to modify this line.
  elif x.ndim in (2, 3):  # Linear layers of shape IO or multihead linear
    squared_norm = jnp.sum(numerics.abs_sq(x), axis=0, keepdims=True)
  elif x.ndim == 4:  # Conv kernels of shape HWIO

# Path: optax/_src/base.py
# NO_PARAMS_MSG = (
#     'You are using a transformation that requires the current value of '
#     'parameters, but you are not passing `params` when calling `update`.')
# class TransformInitFn(typing_extensions.Protocol):
# class TransformUpdateFn(typing_extensions.Protocol):
# class GradientTransformation(NamedTuple):
# class EmptyState(NamedTuple):
#   def __call__(self, params: Params) -> OptState:
#   def __call__(
#       self,
#       updates: Updates,
#       state: OptState,
#       params: Optional[Params] = None
#     ) -> Tuple[Updates, OptState]:
# def identity() -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def set_to_zero() -> GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def stateless(
#     f: Callable[[Updates, Optional[Params]], Updates],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def stateless_with_tree_map(
#     f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# 
# 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/numerics.py
# def abs_sq(x: chex.Array) -> chex.Array:
# def safe_norm(x: chex.Array,
#               min_norm: chex.Numeric,
#               ord: Optional[Union[int, float, str]] = None,  # pylint: disable=redefined-builtin
#               axis: Union[None, Tuple[int, ...], int] = None,
#               keepdims: bool = False) -> chex.Array:
# def safe_root_mean_squares(x: chex.Array, min_rms: chex.Numeric) -> chex.Array:
# def safe_int32_increment(count: chex.Numeric) -> chex.Numeric:


# Path: optax/_src/clipping.py
import chex
import jax
import jax.numpy as jnp
from optax._src import base
from optax._src import linear_algebra
from optax._src import numerics

    del params
    return ClipByGlobalNormState()

  def update_fn(updates, state, params=None):
    del params
    g_norm = linear_algebra.global_norm(updates)
    # TODO(b/163995078): revert back to the following (faster) implementation
    # once analysed how it affects backprop through update (e.g. meta-gradients)
    # g_norm = jnp.maximum(max_norm, g_norm)
    # updates = jax.tree_map(lambda t: (t / g_norm) * max_norm, updates)
    trigger = jnp.squeeze(g_norm < max_norm)
    chex.assert_shape(trigger, ())  # A scalar.

    updates = jax.tree_map(
        lambda t: jax.lax.select(trigger, t, (t / g_norm) * max_norm), updates)
    return updates, state

  return base.GradientTransformation(init_fn, update_fn)


def unitwise_norm(x: chex.Array) -> chex.Array:
  """Computes norms of each output unit separately."""
  if jnp.squeeze(x).ndim <= 1:  # Scalars and vectors
    squared_norm = jnp.sum(numerics.abs_sq(x), keepdims=True)
  # Note that this assumes parameters with a shape of length 3 are multihead
  # linear parameters--if you wish to apply AGC to 1D convs, you may need
  # to modify this line.
  elif x.ndim in (2, 3):  # Linear layers of shape IO or multihead linear
    squared_norm = jnp.sum(numerics.abs_sq(x), axis=0, keepdims=True)
  elif x.ndim == 4:  # Conv kernels of shape HWIO

    squared_norm = jnp.sum(numerics.abs_sq(x), axis=(0, 1, 2), keepdims=True)