# 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):


import chex
import jax
import jax.numpy as jnp
from optax._src import base


def incremental_update(
    new_tensors: base.Params,
    old_tensors: base.Params,
    step_size: chex.Numeric
) -> base.Params:
  """Incrementally update parameters via polyak averaging.

  Polyak averaging tracks an (exponential moving) average of the past
  parameters of a model, for use at test/evaluation time.

  References:
    [Polyak et al, 1991](https://epubs.siam.org/doi/10.1137/0330046)

  Args:
    new_tensors: the latest value of the tensors.
    old_tensors: a moving average of the values of the tensors.
    step_size: the step_size used to update the polyak average on each step.

  Returns:
    an updated moving average `step_size*new+(1-step_size)*old` of the params.
  """
  return jax.tree_multimap(
      lambda new, old: step_size * new + (1.0 - step_size) * old,
      new_tensors, old_tensors)


def periodic_update(
    new_tensors: base.Params,

# 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/update.py
import chex
import jax
import jax.numpy as jnp
from optax._src import base



def incremental_update(
    new_tensors: base.Params,
    old_tensors: base.Params,
    step_size: chex.Numeric
) -> base.Params:
  """Incrementally update parameters via polyak averaging.

  Polyak averaging tracks an (exponential moving) average of the past
  parameters of a model, for use at test/evaluation time.

  References:
    [Polyak et al, 1991](https://epubs.siam.org/doi/10.1137/0330046)

  Args:
    new_tensors: the latest value of the tensors.
    old_tensors: a moving average of the values of the tensors.
    step_size: the step_size used to update the polyak average on each step.

  Returns:
    an updated moving average `step_size*new+(1-step_size)*old` of the params.
  """
  return jax.tree_multimap(
      lambda new, old: step_size * new + (1.0 - step_size) * old,
      new_tensors, old_tensors)


def periodic_update(
    new_tensors: base.Params,

    old_tensors: base.Params,