# 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/types.py



import dataclasses
from typing import Any, Callable, Dict, List, Optional
from flax import linen as nn
from flaxformer.architectures.common import param_remapping
from flaxformer.types import Array


@dataclasses.dataclass
class TransparentLayerSequence:
  """Version of LayerSequence that doesn't add pytree keys.

  Normally one should instantiate the layers in a parent module.

  Attributes:
    layers: List of nn.Modules, which should be owned by a parent Flax module.
  """
  layers: List[nn.Module]

  def __call__(self, inputs: Array, *args, **kwargs) -> Array:
    """Applies all Transformer layers to the inputs sequentially.

    Args:
      inputs: The inputs to the first layer <float>[..., seq_len, hidden_size].
        Typically these are the embedded token IDs, combined with embedded
        position IDs (or sinusoidal position encodings) and segment IDs.
      *args: Positional arguments to be passed to each layer.
      **kwargs: Keyword arguments to be passed to each layer.

    Returns:
      The encoded inputs <float>[..., seq_len, hidden_size].
    """
    # Apply all layers and return the output of the last layer.
    return self.apply_range_of_layers(0, None, inputs, *args, **kwargs)

  def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],

# 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/types.py



# Path: flaxformer/transformer_common.py
import dataclasses
from typing import Any, Callable, Dict, List, Optional
from flax import linen as nn
from flaxformer.architectures.common import param_remapping
from flaxformer.types import Array



@dataclasses.dataclass
class TransparentLayerSequence:
  """Version of LayerSequence that doesn't add pytree keys.

  Normally one should instantiate the layers in a parent module.

  Attributes:
    layers: List of nn.Modules, which should be owned by a parent Flax module.
  """
  layers: List[nn.Module]

  def __call__(self, inputs: Array, *args, **kwargs) -> Array:
    """Applies all Transformer layers to the inputs sequentially.

    Args:
      inputs: The inputs to the first layer <float>[..., seq_len, hidden_size].
        Typically these are the embedded token IDs, combined with embedded
        position IDs (or sinusoidal position encodings) and segment IDs.
      *args: Positional arguments to be passed to each layer.
      **kwargs: Keyword arguments to be passed to each layer.

    Returns:
      The encoded inputs <float>[..., seq_len, hidden_size].
    """
    # Apply all layers and return the output of the last layer.
    return self.apply_range_of_layers(0, None, inputs, *args, **kwargs)

  def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],

                            inputs: Array, *args, **kwargs) -> Array: