# Path: init2winit/optimizer_lib/optmaximus.py
# def from_hparams(opt_hparams):
#   """Create kitchen_sink optimizer from init2winit."""
#   elements = []
#   hps = []
#   masks = []
#   index = 0
# 
#   while str(index) in opt_hparams:
#     hparams = opt_hparams[str(index)]
#     element = hparams['element']
#     hp = hparams.get('hps', {})
#     mask = hparams.get('mask')
# 
#     # TODO(dsuo): there is some unfortunate badness here where `mask` for
#     # `add_decayed_weights` can show up in either the `hps` or in `mask` keys.
#     # If we have less general requirements for masking than we thought, then
#     # we should consider amending the `kitchen_sink` mask API.
#     if element == 'add_decayed_weights' and (hp.get('mask') == 'bias_bn' or
#                                              mask == 'bias_bn'):
#       mask = create_weight_decay_mask()
#       if 'mask' in hp:
#         del hp['mask']
# 
#     for h in copy.deepcopy(hp).keys():
#       if 'one_minus_' in h:
#         hp[h.replace('one_minus_', '')] = 1 - hp[h]
#         del hp[h]
# 
#     elements.append(element)
#     hps.append(hp)
#     masks.append(mask)
#     index += 1
# 
#   return optax.inject_hyperparams(kitchen_sink)(
#       learning_rate=-1.0, elements=elements, hps=hps, masks=masks)
# 
# 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 absl.testing import absltest
from init2winit.optimizer_lib.optmaximus import from_hparams
from init2winit.optimizer_lib.optmaximus import kitchen_sink
import chex
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import ml_collections
import optax
    state = optimizer.init(params)
    update, state = optimizer.update(params, state, params)

    chex.assert_trees_all_close(params, update)

  def test_mask_dim_1(self):
    """Test mask dimension."""
    optimizer = kitchen_sink(
        ['nesterov'], masks=[lambda p: jax.tree_map(lambda x: x.ndim != 1, p)])
    params = {'w': jnp.array([1, 2, 3]), 'b': jnp.ones((2, 2))}
    state = optimizer.init(params)
    update, state = optimizer.update(params, state, params)

    optimizer2 = kitchen_sink(['nesterov'])
    params2 = {'b': jnp.ones((2, 2))}
    state2 = optimizer2.init(params2)
    update2, state2 = optimizer2.update(params2, state2, params2)

    chex.assert_trees_all_close(update['b'], update2['b'])
    chex.assert_trees_all_close(update['w'], params['w'])


class FromHParamsTest(chex.TestCase):
  """Test construction from opt_hparams ml_collections.ConfigDict."""

  def test_empty_element(self):
    self.assertRaises(KeyError, from_hparams,
                      ml_collections.ConfigDict({'0': {}}))

  def test_empty_hps(self):

# Path: init2winit/optimizer_lib/optmaximus.py
# def from_hparams(opt_hparams):
#   """Create kitchen_sink optimizer from init2winit."""
#   elements = []
#   hps = []
#   masks = []
#   index = 0
# 
#   while str(index) in opt_hparams:
#     hparams = opt_hparams[str(index)]
#     element = hparams['element']
#     hp = hparams.get('hps', {})
#     mask = hparams.get('mask')
# 
#     # TODO(dsuo): there is some unfortunate badness here where `mask` for
#     # `add_decayed_weights` can show up in either the `hps` or in `mask` keys.
#     # If we have less general requirements for masking than we thought, then
#     # we should consider amending the `kitchen_sink` mask API.
#     if element == 'add_decayed_weights' and (hp.get('mask') == 'bias_bn' or
#                                              mask == 'bias_bn'):
#       mask = create_weight_decay_mask()
#       if 'mask' in hp:
#         del hp['mask']
# 
#     for h in copy.deepcopy(hp).keys():
#       if 'one_minus_' in h:
#         hp[h.replace('one_minus_', '')] = 1 - hp[h]
#         del hp[h]
# 
#     elements.append(element)
#     hps.append(hp)
#     masks.append(mask)
#     index += 1
# 
#   return optax.inject_hyperparams(kitchen_sink)(
#       learning_rate=-1.0, elements=elements, hps=hps, masks=masks)
# 
# 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_optmaximus.py
from absl.testing import absltest
from init2winit.optimizer_lib.optmaximus import from_hparams
from init2winit.optimizer_lib.optmaximus import kitchen_sink
import chex
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import ml_collections
import optax

    state = optimizer.init(params)
    update, state = optimizer.update(params, state, params)

    chex.assert_trees_all_close(params, update)

  def test_mask_dim_1(self):
    """Test mask dimension."""
    optimizer = kitchen_sink(
        ['nesterov'], masks=[lambda p: jax.tree_map(lambda x: x.ndim != 1, p)])
    params = {'w': jnp.array([1, 2, 3]), 'b': jnp.ones((2, 2))}
    state = optimizer.init(params)
    update, state = optimizer.update(params, state, params)

    optimizer2 = kitchen_sink(['nesterov'])
    params2 = {'b': jnp.ones((2, 2))}
    state2 = optimizer2.init(params2)
    update2, state2 = optimizer2.update(params2, state2, params2)

    chex.assert_trees_all_close(update['b'], update2['b'])
    chex.assert_trees_all_close(update['w'], params['w'])


class FromHParamsTest(chex.TestCase):
  """Test construction from opt_hparams ml_collections.ConfigDict."""

  def test_empty_element(self):
    self.assertRaises(KeyError, from_hparams,
                      ml_collections.ConfigDict({'0': {}}))

  def test_empty_hps(self):

    from_hparams(ml_collections.ConfigDict({'0': {'element': 'nesterov'}}))