# 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
  layers: List[List[Union[OldDense, NewDense]]]

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


class ParameterRemappableTest(parameterized.TestCase):

  @parameterized.named_parameters([
      ('direct_init', Mlp([NewDense(shape=(3, 4)),
                           NewDense(shape=(5, 3))])),
      ('setup_init', MlpSetupInit()),
  ])
  def test_load_old_checkpoint(self, new_mlp):
    # Instantiate an old model and get its params in order to simulate the
    # existence of an old checkpoint.
    old_mlp = Mlp([OldDense(shape=(3, 4)), OldDense(shape=(5, 3))])
    old_mlp_vars = old_mlp.init(jax.random.PRNGKey(0), jnp.zeros([4]))
    self.assertSameStructure(
        testing_utils.param_shapes(old_mlp_vars),
        {'params': {
            'layers_0': {
                'w': [3, 4]
            },
            'layers_1': {
                'w': [5, 3]
            },

# 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

  layers: List[List[Union[OldDense, NewDense]]]

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


class ParameterRemappableTest(parameterized.TestCase):

  @parameterized.named_parameters([
      ('direct_init', Mlp([NewDense(shape=(3, 4)),
                           NewDense(shape=(5, 3))])),
      ('setup_init', MlpSetupInit()),
  ])
  def test_load_old_checkpoint(self, new_mlp):
    # Instantiate an old model and get its params in order to simulate the
    # existence of an old checkpoint.
    old_mlp = Mlp([OldDense(shape=(3, 4)), OldDense(shape=(5, 3))])
    old_mlp_vars = old_mlp.init(jax.random.PRNGKey(0), jnp.zeros([4]))
    self.assertSameStructure(
        testing_utils.param_shapes(old_mlp_vars),
        {'params': {
            'layers_0': {
                'w': [3, 4]
            },
            'layers_1': {
                'w': [5, 3]
            },

        }}, 'old_mlp_vars = ' + testing_utils.format_params_shapes(