# 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
# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for param_remapping."""





class OldDense(nn.Module):
  shape: Tuple[int, ...]

  def setup(self):
    self.w = self.param('w', nn.initializers.normal(), self.shape)

  def __call__(self, x):
    return self.w @ 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

# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for param_remapping."""





class OldDense(nn.Module):
  shape: Tuple[int, ...]

  def setup(self):
    self.w = self.param('w', nn.initializers.normal(), self.shape)

  def __call__(self, x):
    return self.w @ x



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