# Path: flaxformer/testing_utils.py
# def param_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes_axes(params: Mapping[str, Any],
#                              params_axes: Mapping[str, Any]) -> Dict[str, Any]:
#   def remove_axes_suffix(ks):
#   def _create_entry(param, param_axes):
# def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
#   def re_compact_arrays(array_match) -> str:
#   def __init__(self, base_path: str):
#   def get_params(self, filename: str) -> Mapping[str, Any]:
#   def check_params(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_and_axes(
#       self,
#       actual_params: Mapping[str, Any],
#       actual_params_axes: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_shapes_only(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
# class ExpectedJsonFiles:
# 
# Path: flaxformer/architectures/common/param_remapping.py
# METADATA_KEY = '__save_format_metadata__'
# VERSION_KEY = '__version__'
# class ParameterRemappable(metaclass=abc.ABCMeta):
# class RecursiveDefaultDict(collections.defaultdict):
#   def save_format_metadata(self) -> Dict[str, Any]:
#   def from_save_format(self, params: Mapping[str, Any]) -> Dict[str, Any]:
#   def to_save_format(self, params: Mapping[str, Any]) -> Dict[str, Any]:
#   def _from_save_format(self,
#                         params: RecursiveDefaultDict) -> Mapping[str, Any]:
#   def _to_save_format(self, params: RecursiveDefaultDict) -> Mapping[str, Any]:
#   def _submodules_from_save_format(
#       self, params: Mapping[str, Any]) -> RecursiveDefaultDict:
#   def _submodules_to_save_format(
#       self, params: Mapping[str, Any]) -> RecursiveDefaultDict:
#   def _get_remappable_submodules(self) -> List[Tuple[str, ParameterRemappable]]:
#   def _get_submodules(self) -> List[nn.Module]:
#   def __init__(self, initial_content: Optional[Mapping[Any, Any]] = None):
#   def merge(self, key: Any, value: Any) -> None:
#   def update(self, other: Mapping[Any, Any]) -> None:
#   def pop(self, key: Any, default: Any = None) -> Any:
#   def to_dict(self) -> Dict[Any, Any]:
# def filter_out_metadata(params: Mapping[str, Any]) -> Dict[str, Any]:


from typing import Any, List, Mapping, Tuple, Union
from absl.testing import absltest
from absl.testing import parameterized
from flax import linen as nn
from jax import numpy as jnp
from flaxformer import testing_utils
from flaxformer.architectures.common import param_remapping
import flax
import jax
    return params

  @nn.nowrap  # exempt from named call decorator.
  @flax.linen.module.wrap_method_once
  def _to_save_format(
      self, params: param_remapping.RecursiveDefaultDict) -> Mapping[str, Any]:
    params.merge('w', params.pop('weights'))
    return params


class Mlp(nn.Module, param_remapping.ParameterRemappable):
  layers: List[Union[OldDense, NewDense]]

  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    return x


class MlpSetupInit(nn.Module, param_remapping.ParameterRemappable):

  def setup(self):
    self.layers = [NewDense(shape=(3, 4)), NewDense(shape=(5, 3))]

  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    return x



# Path: flaxformer/testing_utils.py
# def param_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes_axes(params: Mapping[str, Any],
#                              params_axes: Mapping[str, Any]) -> Dict[str, Any]:
#   def remove_axes_suffix(ks):
#   def _create_entry(param, param_axes):
# def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
#   def re_compact_arrays(array_match) -> str:
#   def __init__(self, base_path: str):
#   def get_params(self, filename: str) -> Mapping[str, Any]:
#   def check_params(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_and_axes(
#       self,
#       actual_params: Mapping[str, Any],
#       actual_params_axes: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_shapes_only(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
# class ExpectedJsonFiles:
# 
# Path: flaxformer/architectures/common/param_remapping.py
# METADATA_KEY = '__save_format_metadata__'
# VERSION_KEY = '__version__'
# class ParameterRemappable(metaclass=abc.ABCMeta):
# class RecursiveDefaultDict(collections.defaultdict):
#   def save_format_metadata(self) -> Dict[str, Any]:
#   def from_save_format(self, params: Mapping[str, Any]) -> Dict[str, Any]:
#   def to_save_format(self, params: Mapping[str, Any]) -> Dict[str, Any]:
#   def _from_save_format(self,
#                         params: RecursiveDefaultDict) -> Mapping[str, Any]:
#   def _to_save_format(self, params: RecursiveDefaultDict) -> Mapping[str, Any]:
#   def _submodules_from_save_format(
#       self, params: Mapping[str, Any]) -> RecursiveDefaultDict:
#   def _submodules_to_save_format(
#       self, params: Mapping[str, Any]) -> RecursiveDefaultDict:
#   def _get_remappable_submodules(self) -> List[Tuple[str, ParameterRemappable]]:
#   def _get_submodules(self) -> List[nn.Module]:
#   def __init__(self, initial_content: Optional[Mapping[Any, Any]] = None):
#   def merge(self, key: Any, value: Any) -> None:
#   def update(self, other: Mapping[Any, Any]) -> None:
#   def pop(self, key: Any, default: Any = None) -> Any:
#   def to_dict(self) -> Dict[Any, Any]:
# def filter_out_metadata(params: Mapping[str, Any]) -> Dict[str, Any]:


# Path: flaxformer/architectures/common/param_remapping_test.py
from typing import Any, List, Mapping, Tuple, Union
from absl.testing import absltest
from absl.testing import parameterized
from flax import linen as nn
from jax import numpy as jnp
from flaxformer import testing_utils
from flaxformer.architectures.common import param_remapping
import flax
import jax

    return params

  @nn.nowrap  # exempt from named call decorator.
  @flax.linen.module.wrap_method_once
  def _to_save_format(
      self, params: param_remapping.RecursiveDefaultDict) -> Mapping[str, Any]:
    params.merge('w', params.pop('weights'))
    return params


class Mlp(nn.Module, param_remapping.ParameterRemappable):
  layers: List[Union[OldDense, NewDense]]

  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    return x


class MlpSetupInit(nn.Module, param_remapping.ParameterRemappable):

  def setup(self):
    self.layers = [NewDense(shape=(3, 4)), NewDense(shape=(5, 3))]

  def __call__(self, x):
    for layer in self.layers:
      x = layer(x)
    return x



class NestedStructuresMlp(nn.Module, param_remapping.ParameterRemappable):