# 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/wrappers.py
# def flatten(
#     inner: base.GradientTransformation
# ) -> base.GradientTransformation:
#   def _flatten(params):
#   def _unflatten(updates, flat):
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def apply_if_finite(
#     inner: base.GradientTransformation,
#     max_consecutive_errors: int
# ) -> base.GradientTransformation:
#   def init(params):
#   def update(updates, state, params=None):
#     def do_update(_):
#     def reject_update(_):
# def _zeros_tree_like(inp_tree):
#   def __init__(self,
#                opt: base.GradientTransformation,
#                every_k_schedule: Union[int, Callable[[Array], Array]],
#                use_grad_mean: bool = True):
#   def inner_opt(self):
#   def init(self, params: Any) -> MultiStepsState:
#   def update(self,
#              updates: base.Updates,
#              state: MultiStepsState,
#              params: Optional[base.Params] = None
#              ) -> Tuple[base.Updates, MultiStepsState]:
#     def final_step(args):
#     def mid_step(args):
#   def has_updated(self, state: MultiStepsState) -> Array:
#   def gradient_transformation(self) -> base.GradientTransformation:
# def masked(
#     inner: base.GradientTransformation,
#     mask: Union[base.PyTree, Callable[[base.Params], base.PyTree]]
# ) -> base.GradientTransformation:
#   def mask_pytree(pytree, mask_tree):
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def maybe_update(
#     inner: base.GradientTransformation,
#     should_update_fn: Callable[[Array], Array]
# ) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
#     def do_update(_):
#     def reject_update(_):
# class ApplyIfFiniteState(NamedTuple):
# class MultiStepsState(NamedTuple):
# class MultiSteps:
# class MaskedState(NamedTuple):
# class MaybeUpdateState(NamedTuple):


from typing import Callable, NamedTuple, Union, Mapping, Hashable
from optax._src import base
from optax._src import wrappers
import jax

  Returns:
    A single (init_fn, update_fn) tuple.
  """

  init_fns, update_fns = zip(*args)

  def init_fn(params):
    return tuple(fn(params) for fn in init_fns)

  def update_fn(updates, state, params=None):
    if len(update_fns) != len(state):
      raise ValueError('The number of updates and states has to be the same in '
                       'chain! Make sure you have called init first!')

    new_state = []
    for s, fn in zip(state, update_fns):
      updates, new_s = fn(updates, s, params)
      new_state.append(new_s)
    return updates, tuple(new_state)

  return base.GradientTransformation(init_fn, update_fn)


class MultiTransformState(NamedTuple):
  inner_states: Mapping[Hashable, NamedTuple]


def multi_transform(
    transforms: Mapping[Hashable, base.GradientTransformation],

# 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/wrappers.py
# def flatten(
#     inner: base.GradientTransformation
# ) -> base.GradientTransformation:
#   def _flatten(params):
#   def _unflatten(updates, flat):
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def apply_if_finite(
#     inner: base.GradientTransformation,
#     max_consecutive_errors: int
# ) -> base.GradientTransformation:
#   def init(params):
#   def update(updates, state, params=None):
#     def do_update(_):
#     def reject_update(_):
# def _zeros_tree_like(inp_tree):
#   def __init__(self,
#                opt: base.GradientTransformation,
#                every_k_schedule: Union[int, Callable[[Array], Array]],
#                use_grad_mean: bool = True):
#   def inner_opt(self):
#   def init(self, params: Any) -> MultiStepsState:
#   def update(self,
#              updates: base.Updates,
#              state: MultiStepsState,
#              params: Optional[base.Params] = None
#              ) -> Tuple[base.Updates, MultiStepsState]:
#     def final_step(args):
#     def mid_step(args):
#   def has_updated(self, state: MultiStepsState) -> Array:
#   def gradient_transformation(self) -> base.GradientTransformation:
# def masked(
#     inner: base.GradientTransformation,
#     mask: Union[base.PyTree, Callable[[base.Params], base.PyTree]]
# ) -> base.GradientTransformation:
#   def mask_pytree(pytree, mask_tree):
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def maybe_update(
#     inner: base.GradientTransformation,
#     should_update_fn: Callable[[Array], Array]
# ) -> base.GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
#     def do_update(_):
#     def reject_update(_):
# class ApplyIfFiniteState(NamedTuple):
# class MultiStepsState(NamedTuple):
# class MultiSteps:
# class MaskedState(NamedTuple):
# class MaybeUpdateState(NamedTuple):


# Path: optax/_src/combine.py
from typing import Callable, NamedTuple, Union, Mapping, Hashable
from optax._src import base
from optax._src import wrappers
import jax


  Returns:
    A single (init_fn, update_fn) tuple.
  """

  init_fns, update_fns = zip(*args)

  def init_fn(params):
    return tuple(fn(params) for fn in init_fns)

  def update_fn(updates, state, params=None):
    if len(update_fns) != len(state):
      raise ValueError('The number of updates and states has to be the same in '
                       'chain! Make sure you have called init first!')

    new_state = []
    for s, fn in zip(state, update_fns):
      updates, new_s = fn(updates, s, params)
      new_state.append(new_s)
    return updates, tuple(new_state)

  return base.GradientTransformation(init_fn, update_fn)


class MultiTransformState(NamedTuple):
  inner_states: Mapping[Hashable, NamedTuple]


def multi_transform(
    transforms: Mapping[Hashable, base.GradientTransformation],

    param_labels: Union[base.PyTree, Callable[[base.PyTree], base.PyTree]]