# 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, Callable, Dict, Mapping, Sequence, Tuple
from absl import logging
from flax import linen as nn
from flax import traverse_util
from flax.core import frozen_dict
from t5x import checkpoints
from flaxformer.architectures.common import param_remapping



# frozen_dict.unfreeze is incorrectly typed, so introduce an alias.
def _unfreeze(x: Mapping[str, Any]) -> Dict[str, Any]:
  return frozen_dict.unfreeze(x)  # pytype: disable=wrong-arg-types


def _flattened_names(state_dict: Mapping[str, Any]) -> Sequence[str]:
  return [
      '/'.join(k)
      for k in traverse_util.flatten_dict(state_dict, keep_empty_nodes=True)
  ]


def _apply_remap_fn(remap_fn: Callable[[Mapping[str, Any]], Mapping[str, Any]],
                    state_dict: Mapping[str, Any]) -> Mapping[str, Any]:
  result = _unfreeze(state_dict)
  result['state']['param_states'] = remap_fn(result['state']['param_states'])
  result['target'] = remap_fn(result['target'])
  return result


def make_to_save_format_fn(
    module: nn.Module) -> checkpoints.SaveStateTransformationFn:
  """Returns a t5x on-save state transformation function.

  Args:
    module: A Flax module inheriting from param_remapping.ParamRemappable.
  """

# 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/t5x/checkpointing_util.py
from typing import Any, Callable, Dict, Mapping, Sequence, Tuple
from absl import logging
from flax import linen as nn
from flax import traverse_util
from flax.core import frozen_dict
from t5x import checkpoints
from flaxformer.architectures.common import param_remapping




# frozen_dict.unfreeze is incorrectly typed, so introduce an alias.
def _unfreeze(x: Mapping[str, Any]) -> Dict[str, Any]:
  return frozen_dict.unfreeze(x)  # pytype: disable=wrong-arg-types


def _flattened_names(state_dict: Mapping[str, Any]) -> Sequence[str]:
  return [
      '/'.join(k)
      for k in traverse_util.flatten_dict(state_dict, keep_empty_nodes=True)
  ]


def _apply_remap_fn(remap_fn: Callable[[Mapping[str, Any]], Mapping[str, Any]],
                    state_dict: Mapping[str, Any]) -> Mapping[str, Any]:
  result = _unfreeze(state_dict)
  result['state']['param_states'] = remap_fn(result['state']['param_states'])
  result['target'] = remap_fn(result['target'])
  return result


def make_to_save_format_fn(
    module: nn.Module) -> checkpoints.SaveStateTransformationFn:
  """Returns a t5x on-save state transformation function.

  Args:
    module: A Flax module inheriting from param_remapping.ParamRemappable.
  """

  if not isinstance(module, param_remapping.ParameterRemappable):