# 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
#
# 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.

"""Common Transformer classes."""




class LayerSequence(nn.Module, param_remapping.ParameterRemappable):
  """Common object responsible for holding and applying Transformer layers.

  Can be used in encoders/decoders and supports applying a range of sublayers.

  Attributes:
    num_layers: The number of Transformer layers to create for this encoder.
    make_layer: A function that returns a single encoder layer.
    extra_modules: Extra modules provided to each make_layer call as kwargs.
  """
  num_layers: int
  make_layer: Callable[..., Any]
  extra_modules: Optional[Dict[str, Optional[nn.Module]]] = None

  def setup(self):
    kwargs = self.extra_modules if self.extra_modules is not None else {}
    self.layers = [self.make_layer(**kwargs) for _ in range(self.num_layers)]


# 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

#
# 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.

"""Common Transformer classes."""




class LayerSequence(nn.Module, param_remapping.ParameterRemappable):
  """Common object responsible for holding and applying Transformer layers.

  Can be used in encoders/decoders and supports applying a range of sublayers.

  Attributes:
    num_layers: The number of Transformer layers to create for this encoder.
    make_layer: A function that returns a single encoder layer.
    extra_modules: Extra modules provided to each make_layer call as kwargs.
  """
  num_layers: int
  make_layer: Callable[..., Any]
  extra_modules: Optional[Dict[str, Optional[nn.Module]]] = None

  def setup(self):
    kwargs = self.extra_modules if self.extra_modules is not None else {}
    self.layers = [self.make_layer(**kwargs) for _ in range(self.num_layers)]


  def __call__(self, inputs: Array, *args, **kwargs) -> Array: