# 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
      logging.info('to_save_format input state_dict: %s', name)
    for name in _flattened_names(parameter_infos):
      logging.info('to_save_format input parameter_infos: %s', name)

    result_state_dict = _apply_remap_fn(remap, state_dict)
    result_parameter_infos = _apply_remap_fn(remap, parameter_infos)

    flat_parameter_infos = traverse_util.flatten_dict(
        result_parameter_infos, keep_empty_nodes=True)
    result_parameter_infos = traverse_util.unflatten_dict({
        k: (v if k[-1] != param_remapping.VERSION_KEY else None)
        for k, v in flat_parameter_infos.items()
    })

    for name in _flattened_names(result_state_dict):
      logging.info('to_save_format output state_dict: %s', name)
    for name in _flattened_names(result_parameter_infos):
      logging.info('to_save_format output parameter_infos: %s', name)
    return result_state_dict, result_parameter_infos

  return to_save_format


def make_from_save_format_fn(
    module: nn.Module) -> checkpoints.RestoreStateTransformationFn:
  """Returns a t5x on-restore 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

      logging.info('to_save_format input state_dict: %s', name)
    for name in _flattened_names(parameter_infos):
      logging.info('to_save_format input parameter_infos: %s', name)

    result_state_dict = _apply_remap_fn(remap, state_dict)
    result_parameter_infos = _apply_remap_fn(remap, parameter_infos)

    flat_parameter_infos = traverse_util.flatten_dict(
        result_parameter_infos, keep_empty_nodes=True)
    result_parameter_infos = traverse_util.unflatten_dict({
        k: (v if k[-1] != param_remapping.VERSION_KEY else None)
        for k, v in flat_parameter_infos.items()
    })

    for name in _flattened_names(result_state_dict):
      logging.info('to_save_format output state_dict: %s', name)
    for name in _flattened_names(result_parameter_infos):
      logging.info('to_save_format output parameter_infos: %s', name)
    return result_state_dict, result_parameter_infos

  return to_save_format


def make_from_save_format_fn(
    module: nn.Module) -> checkpoints.RestoreStateTransformationFn:
  """Returns a t5x on-restore state transformation function.

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

  if not isinstance(module, param_remapping.ParameterRemappable):