# Path: init2winit/optimizer_lib/utils.py
# def create_mask(fn):
#   """Creates a mask that maps fn over the leaves of a dict.
# 
#   Args:
#     fn: function to apply taking
#       - k: Tuple containing nodes (strings) in path to the leaf
#       - v: The leaf
# 
#   Returns:
#     mask: function that takes dict and returns mapped dict
#   """
# 
#   def mask(data):
#     flattened_dict = flax.traverse_util.flatten_dict(data)
#     return flax.traverse_util.unflatten_dict(
#         {k: fn(k, v) for k, v in flattened_dict.items()})
# 
#   return mask
# 
# Path: init2winit/optimizer_lib/utils.py
# def create_weight_decay_mask():
#   return create_mask(
#       lambda p, _: 'bias' not in p and not p[-2].startswith('BatchNorm'))


from typing import Sequence
from absl.testing import absltest
from init2winit.optimizer_lib.utils import create_mask
from init2winit.optimizer_lib.utils import create_weight_decay_mask
import chex
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax
    y = jax.random.uniform(key, (5, 4, 4, 7))

    foo_vars = flax.core.unfreeze(Foo(filters=7, train=True).init(key, x))
    tx = optax.masked(optax.adam(1e-7), create_weight_decay_mask())

    @self.variant
    def train_step(params, x, y):
      y1, new_batch_stats = Foo(
          filters=7, train=True).apply(
              params, x, mutable=['batch_stats'])

      return jnp.abs(y - y1).sum(), new_batch_stats

    state = self.variant(tx.init)(foo_vars['params'])
    grads, _ = jax.grad(train_step, has_aux=True)(foo_vars, x, y)
    updates, state = self.variant(tx.update)(dict(grads['params']), state)

    chex.assert_trees_all_close(updates['BatchNorm_0'],
                                grads['params']['BatchNorm_0'])

  @chex.variants(with_jit=True, without_jit=True)
  def test_bias(self):
    """Test that biases are ignored."""
    key1, key2 = jax.random.split(jax.random.PRNGKey(0), 2)
    x = jax.random.uniform(key1, (4, 4))

    model = Bar(features=[3, 4, 5])
    params = flax.core.unfreeze(model.init(key2, x))
    y = jax.random.uniform(key1, model.apply(params, x).shape)


# Path: init2winit/optimizer_lib/utils.py
# def create_mask(fn):
#   """Creates a mask that maps fn over the leaves of a dict.
# 
#   Args:
#     fn: function to apply taking
#       - k: Tuple containing nodes (strings) in path to the leaf
#       - v: The leaf
# 
#   Returns:
#     mask: function that takes dict and returns mapped dict
#   """
# 
#   def mask(data):
#     flattened_dict = flax.traverse_util.flatten_dict(data)
#     return flax.traverse_util.unflatten_dict(
#         {k: fn(k, v) for k, v in flattened_dict.items()})
# 
#   return mask
# 
# Path: init2winit/optimizer_lib/utils.py
# def create_weight_decay_mask():
#   return create_mask(
#       lambda p, _: 'bias' not in p and not p[-2].startswith('BatchNorm'))


# Path: init2winit/optimizer_lib/test_utils.py
from typing import Sequence
from absl.testing import absltest
from init2winit.optimizer_lib.utils import create_mask
from init2winit.optimizer_lib.utils import create_weight_decay_mask
import chex
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax

    y = jax.random.uniform(key, (5, 4, 4, 7))

    foo_vars = flax.core.unfreeze(Foo(filters=7, train=True).init(key, x))
    tx = optax.masked(optax.adam(1e-7), create_weight_decay_mask())

    @self.variant
    def train_step(params, x, y):
      y1, new_batch_stats = Foo(
          filters=7, train=True).apply(
              params, x, mutable=['batch_stats'])

      return jnp.abs(y - y1).sum(), new_batch_stats

    state = self.variant(tx.init)(foo_vars['params'])
    grads, _ = jax.grad(train_step, has_aux=True)(foo_vars, x, y)
    updates, state = self.variant(tx.update)(dict(grads['params']), state)

    chex.assert_trees_all_close(updates['BatchNorm_0'],
                                grads['params']['BatchNorm_0'])

  @chex.variants(with_jit=True, without_jit=True)
  def test_bias(self):
    """Test that biases are ignored."""
    key1, key2 = jax.random.split(jax.random.PRNGKey(0), 2)
    x = jax.random.uniform(key1, (4, 4))

    model = Bar(features=[3, 4, 5])
    params = flax.core.unfreeze(model.init(key2, x))
    y = jax.random.uniform(key1, model.apply(params, x).shape)


    tx = optax.masked(optax.adam(1e-7), create_weight_decay_mask())