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


from absl.testing import absltest
from optax._src import base
import chex
import jax
import jax.numpy as jnp
import numpy as np
    @base.stateless
    def opt(g, p):
      return jax.tree_map(lambda g_, p_: g_ + 0.1 * p_, g, p)

    state = opt.init(params)
    update_fn = self.variant(opt.update)
    new_updates, _ = update_fn(updates, state, params)
    expected_updates = {'a': jnp.ones((1, 2)), 'b': jnp.array([2.1])}
    chex.assert_trees_all_close(new_updates, expected_updates)

  @chex.all_variants
  def test_stateless_no_params(self):
    updates = {'linear': jnp.full((5, 3), 3.0)}

    @base.stateless
    def opt(g, _):
      return jax.tree_map(lambda g_: g_ * 2, g)

    state = opt.init(None)
    update_fn = self.variant(opt.update)
    new_updates, _ = update_fn(updates, state)
    expected_updates = {'linear': jnp.full((5, 3), 6.0)}
    chex.assert_trees_all_close(new_updates, expected_updates)

  def test_init_returns_emptystate(self):
    def weight_decay(g, p):
      return jax.tree_map(lambda g_, p_: g_ + 0.1 * p_, g, p)

    opt = base.stateless(weight_decay)
    state = opt.init(None)

# 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/base_test.py
from absl.testing import absltest
from optax._src import base
import chex
import jax
import jax.numpy as jnp
import numpy as np

    @base.stateless
    def opt(g, p):
      return jax.tree_map(lambda g_, p_: g_ + 0.1 * p_, g, p)

    state = opt.init(params)
    update_fn = self.variant(opt.update)
    new_updates, _ = update_fn(updates, state, params)
    expected_updates = {'a': jnp.ones((1, 2)), 'b': jnp.array([2.1])}
    chex.assert_trees_all_close(new_updates, expected_updates)

  @chex.all_variants
  def test_stateless_no_params(self):
    updates = {'linear': jnp.full((5, 3), 3.0)}

    @base.stateless
    def opt(g, _):
      return jax.tree_map(lambda g_: g_ * 2, g)

    state = opt.init(None)
    update_fn = self.variant(opt.update)
    new_updates, _ = update_fn(updates, state)
    expected_updates = {'linear': jnp.full((5, 3), 6.0)}
    chex.assert_trees_all_close(new_updates, expected_updates)

  def test_init_returns_emptystate(self):
    def weight_decay(g, p):
      return jax.tree_map(lambda g_, p_: g_ + 0.1 * p_, g, p)

    opt = base.stateless(weight_decay)
    state = opt.init(None)

    self.assertIsInstance(state, base.EmptyState)