# 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
    x = nn.Conv(self.filters, (1, 1), use_bias=False, dtype=jnp.float32)(x)
    x = nn.BatchNorm(
        use_running_average=not self.train,
        momentum=0.9,
        epsilon=1e-5,
        dtype=jnp.float32)(
            x)
    return x


class Bar(nn.Module):
  """Dummy model."""

  features: Sequence[int]

  @nn.compact
  def __call__(self, inputs):
    x = inputs
    for i, feat in enumerate(self.features):
      x = nn.Dense(feat, use_bias=True, name=f'layers_{i}')(x)
      if i != len(self.features) - 1:
        x = nn.relu(x)
    return x


class CreateMaskTest(chex.TestCase):
  """Test masking."""

  def test_simple(self):
    """Check if the leaf key is `a`."""

# 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

    x = nn.Conv(self.filters, (1, 1), use_bias=False, dtype=jnp.float32)(x)
    x = nn.BatchNorm(
        use_running_average=not self.train,
        momentum=0.9,
        epsilon=1e-5,
        dtype=jnp.float32)(
            x)
    return x


class Bar(nn.Module):
  """Dummy model."""

  features: Sequence[int]

  @nn.compact
  def __call__(self, inputs):
    x = inputs
    for i, feat in enumerate(self.features):
      x = nn.Dense(feat, use_bias=True, name=f'layers_{i}')(x)
      if i != len(self.features) - 1:
        x = nn.relu(x)
    return x


class CreateMaskTest(chex.TestCase):
  """Test masking."""

  def test_simple(self):
    """Check if the leaf key is `a`."""

    mask = create_mask(lambda path, _: path[-1] == 'a')