# Path: init2winit/optimizer_lib/optmaximus.py
# def kitchen_sink(elements: List[str],
#                  hps: List[Dict[str, float]] = None,
#                  masks: List[Any] = None,
#                  learning_rate=-1.) -> optax.GradientTransformation:
#   """Utility function for chaining GradientTransforms based on string names."""
# 
#   hps = hps or [{}] * len(elements)
#   masks = masks or [None] * len(elements)
#   transforms = []
# 
#   if len(hps) != len(elements):
#     raise ValueError('Number of hps must equal number of elements.')
# 
#   if len(masks) != len(elements):
#     raise ValueError('Number of masks must equal number of elements.')
# 
#   transforms = [_transformations[el](**hp) for el, hp in zip(elements, hps)]
# 
#   for i, (transform, mask) in enumerate(zip(transforms, masks)):
#     if mask is not None:
#       transforms[i] = optax.masked(transform, mask)
# 
#   transforms += [scale_by_learning_rate(learning_rate)]
# 
#   init_fn, update_fn = optax.chain(*transforms)
# 
#   # NOTE(dsuo): We use plain dicts internally due to this issue
#   # https://github.com/deepmind/optax/issues/160.
#   def wrapped_init_fn(params):
#     return init_fn(flax.core.unfreeze(params))
# 
#   def wrapped_update_fn(updates, state, params=None):
#     new_updates, state = update_fn(
#         flax.core.unfreeze(updates), state,
#         None if params is None else flax.core.unfreeze(params))
# 
#     if isinstance(updates, flax.core.FrozenDict):
#       new_updates = flax.core.freeze(new_updates)
# 
#     return new_updates, state
# 
#   return optax.GradientTransformation(wrapped_init_fn, wrapped_update_fn)


from typing import NamedTuple
from absl.testing import absltest
from init2winit.optimizer_lib.optmaximus import kitchen_sink
import chex
import jax
import jax.numpy as jnp
import optax
      chex.assert_trees_all_close(target, result)


class EquivalenceTest(chex.TestCase):
  """Test equivalence of kitchen_sink and optax adagrad."""

  def test_adagrad(self):
    true_adagrad = optax.adagrad(0.7, initial_accumulator_value=0.3)
    ks_adagrad = kitchen_sink(['precondition_by_rss', 'first_moment_ema'], [{
        'initial_accumulator_value': 0.3
    }, {
        'decay': 0.0
    }],
                              learning_rate=0.7)

    targets = _optimizer_loop(true_adagrad)
    results = _optimizer_loop(ks_adagrad)

    for target, result in zip(targets, results):
      chex.assert_trees_all_close(target, result)


if __name__ == '__main__':
  absltest.main()


class EqEMAHBTest(chex.TestCase):
  """Tests the equivalence of EMA vs HB. Both LR and WD need to scale."""

  def test_equivalence(self):

# Path: init2winit/optimizer_lib/optmaximus.py
# def kitchen_sink(elements: List[str],
#                  hps: List[Dict[str, float]] = None,
#                  masks: List[Any] = None,
#                  learning_rate=-1.) -> optax.GradientTransformation:
#   """Utility function for chaining GradientTransforms based on string names."""
# 
#   hps = hps or [{}] * len(elements)
#   masks = masks or [None] * len(elements)
#   transforms = []
# 
#   if len(hps) != len(elements):
#     raise ValueError('Number of hps must equal number of elements.')
# 
#   if len(masks) != len(elements):
#     raise ValueError('Number of masks must equal number of elements.')
# 
#   transforms = [_transformations[el](**hp) for el, hp in zip(elements, hps)]
# 
#   for i, (transform, mask) in enumerate(zip(transforms, masks)):
#     if mask is not None:
#       transforms[i] = optax.masked(transform, mask)
# 
#   transforms += [scale_by_learning_rate(learning_rate)]
# 
#   init_fn, update_fn = optax.chain(*transforms)
# 
#   # NOTE(dsuo): We use plain dicts internally due to this issue
#   # https://github.com/deepmind/optax/issues/160.
#   def wrapped_init_fn(params):
#     return init_fn(flax.core.unfreeze(params))
# 
#   def wrapped_update_fn(updates, state, params=None):
#     new_updates, state = update_fn(
#         flax.core.unfreeze(updates), state,
#         None if params is None else flax.core.unfreeze(params))
# 
#     if isinstance(updates, flax.core.FrozenDict):
#       new_updates = flax.core.freeze(new_updates)
# 
#     return new_updates, state
# 
#   return optax.GradientTransformation(wrapped_init_fn, wrapped_update_fn)


# Path: init2winit/optimizer_lib/test_transform.py
from typing import NamedTuple
from absl.testing import absltest
from init2winit.optimizer_lib.optmaximus import kitchen_sink
import chex
import jax
import jax.numpy as jnp
import optax

      chex.assert_trees_all_close(target, result)


class EquivalenceTest(chex.TestCase):
  """Test equivalence of kitchen_sink and optax adagrad."""

  def test_adagrad(self):
    true_adagrad = optax.adagrad(0.7, initial_accumulator_value=0.3)
    ks_adagrad = kitchen_sink(['precondition_by_rss', 'first_moment_ema'], [{
        'initial_accumulator_value': 0.3
    }, {
        'decay': 0.0
    }],
                              learning_rate=0.7)

    targets = _optimizer_loop(true_adagrad)
    results = _optimizer_loop(ks_adagrad)

    for target, result in zip(targets, results):
      chex.assert_trees_all_close(target, result)


if __name__ == '__main__':
  absltest.main()


class EqEMAHBTest(chex.TestCase):
  """Tests the equivalence of EMA vs HB. Both LR and WD need to scale."""

  def test_equivalence(self):

    hb = kitchen_sink(