# 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
    grads = jax.grad(compute_loss)(params, jnp.array([5.0, 6.0]), 4.0)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    results.append(params)
  return results


class NesterovTest(chex.TestCase):
  """Test correctness of nesterov momentum."""

  def test_correctness(self):
    """Testing correctness via an independent flax.optim run."""

    target_solution = [
        {
            'w': jnp.array([0.40500003, 0.286])
        },
        {
            'w': jnp.array([0.255515, 0.106618])
        },
        {
            'w': jnp.array([0.31884143, 0.18260972])
        },
        {
            'w': jnp.array([0.40163627, 0.28196353])
        },
        {
            'w': jnp.array([0.43924114, 0.32708937])
        },
    ]

# 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

    grads = jax.grad(compute_loss)(params, jnp.array([5.0, 6.0]), 4.0)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    results.append(params)
  return results


class NesterovTest(chex.TestCase):
  """Test correctness of nesterov momentum."""

  def test_correctness(self):
    """Testing correctness via an independent flax.optim run."""

    target_solution = [
        {
            'w': jnp.array([0.40500003, 0.286])
        },
        {
            'w': jnp.array([0.255515, 0.106618])
        },
        {
            'w': jnp.array([0.31884143, 0.18260972])
        },
        {
            'w': jnp.array([0.40163627, 0.28196353])
        },
        {
            'w': jnp.array([0.43924114, 0.32708937])
        },
    ]

    optimizer = kitchen_sink(['nesterov'], [{'decay': 0.7}], learning_rate=0.01)