# 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 typing import Any, NamedTuple
from optax._src import base
import jax
import jax.numpy as jnp

  Returns:
    An (init_fn, update_fn) tuple.
  """

  def init_fn(params):
    del params
    return NonNegativeParamsState()

  def update_fn(updates, state, params):
    if params is None:
      raise ValueError(base.NO_PARAMS_MSG)

    updates = jax.tree_multimap(
        lambda p, u: jnp.where((p + u) < 0., -p, u), params, updates)
    return updates, state

  return base.GradientTransformation(init_fn, update_fn)


class ZeroNansState(NamedTuple):
  """Contains a tree.

  The entry `found_nan` has the same tree structure as that of the parameters.
  Each leaf is a single boolean which contains True iff a NaN was detected in
  the corresponding parameter array at the last call to `update`.
  """
  found_nan: Any



# 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/constrain.py
from typing import Any, NamedTuple
from optax._src import base
import jax
import jax.numpy as jnp


  Returns:
    An (init_fn, update_fn) tuple.
  """

  def init_fn(params):
    del params
    return NonNegativeParamsState()

  def update_fn(updates, state, params):
    if params is None:
      raise ValueError(base.NO_PARAMS_MSG)

    updates = jax.tree_multimap(
        lambda p, u: jnp.where((p + u) < 0., -p, u), params, updates)
    return updates, state

  return base.GradientTransformation(init_fn, update_fn)


class ZeroNansState(NamedTuple):
  """Contains a tree.

  The entry `found_nan` has the same tree structure as that of the parameters.
  Each leaf is a single boolean which contains True iff a NaN was detected in
  the corresponding parameter array at the last call to `update`.
  """
  found_nan: Any



def zero_nans() -> base.GradientTransformation: