# 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/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 functools
import jax
import jax.numpy as jnp
import numpy as np
from typing import Any, Callable, NamedTuple, Optional, Tuple, Union
from jax import lax
from jax.tree_util import tree_flatten
from jax.tree_util import tree_map
from jax.tree_util import tree_unflatten
from optax._src import base
from optax._src import numerics

  def update(self,
             updates: base.Updates,
             state: MultiStepsState,
             params: Optional[base.Params] = None
             ) -> Tuple[base.Updates, MultiStepsState]:
    """Accumulates gradients and proposes non-zero updates every `k_steps`."""
    k_steps = self._every_k_schedule(state.gradient_step)
    acc_grads = jax.tree_util.tree_multimap(
        functools.partial(self._acc_update, n_acc=state.mini_step),
        updates, state.acc_grads)

    def final_step(args):
      del args
      updates, new_inner_state = self._opt.update(
          acc_grads, state.inner_opt_state, params=params)
      new_state = MultiStepsState(
          mini_step=jnp.zeros([], dtype=jnp.int32),
          gradient_step=numerics.safe_int32_increment(state.gradient_step),
          inner_opt_state=new_inner_state,
          acc_grads=_zeros_tree_like(acc_grads))
      return updates, new_state

    def mid_step(args):
      del args
      updates_shape_dtype, _ = jax.eval_shape(
          self._opt.update, acc_grads, state.inner_opt_state, params=params)
      updates = jax.tree_map(lambda sd: jnp.zeros(sd.shape, sd.dtype),
                             updates_shape_dtype)
      new_state = MultiStepsState(

# 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/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/wrappers.py
import functools
import jax
import jax.numpy as jnp
import numpy as np
from typing import Any, Callable, NamedTuple, Optional, Tuple, Union
from jax import lax
from jax.tree_util import tree_flatten
from jax.tree_util import tree_map
from jax.tree_util import tree_unflatten
from optax._src import base
from optax._src import numerics


  def update(self,
             updates: base.Updates,
             state: MultiStepsState,
             params: Optional[base.Params] = None
             ) -> Tuple[base.Updates, MultiStepsState]:
    """Accumulates gradients and proposes non-zero updates every `k_steps`."""
    k_steps = self._every_k_schedule(state.gradient_step)
    acc_grads = jax.tree_util.tree_multimap(
        functools.partial(self._acc_update, n_acc=state.mini_step),
        updates, state.acc_grads)

    def final_step(args):
      del args
      updates, new_inner_state = self._opt.update(
          acc_grads, state.inner_opt_state, params=params)
      new_state = MultiStepsState(
          mini_step=jnp.zeros([], dtype=jnp.int32),
          gradient_step=numerics.safe_int32_increment(state.gradient_step),
          inner_opt_state=new_inner_state,
          acc_grads=_zeros_tree_like(acc_grads))
      return updates, new_state

    def mid_step(args):
      del args
      updates_shape_dtype, _ = jax.eval_shape(
          self._opt.update, acc_grads, state.inner_opt_state, params=params)
      updates = jax.tree_map(lambda sd: jnp.zeros(sd.shape, sd.dtype),
                             updates_shape_dtype)
      new_state = MultiStepsState(

          mini_step=numerics.safe_int32_increment(state.mini_step),