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

# 
# Path: flaxformer/types.py



from typing import Optional
from flax import linen as nn
from flax.linen import partitioning
from jax import lax
from jax import numpy as jnp
from flaxformer.architectures.common import param_remapping
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer

class T5LayerNorm(nn.Module, param_remapping.ParameterRemappable):
  """T5 Layer normalization.

  Operates on the last axis of the input data.

  Attributes:
    epsilon: A small float added to variance to avoid dividing by zero.
    dtype: the dtype of the computation (default: float32).
    scale_init: Initializer for scale, by default, one.
    use_scale: boolean - whether to scale by a learned per-channel value
    conv: optional convolution to happen after layer norm
    center_scale_at_zero: boolean - If True, the scale weights will be
      initialized to 0.0 and a constant 1.0 will be added before they are
      applied. Preferable if you use weight decay in the optimizer.
    scale_axis_name: Axis name of the scale variable.
  """
  epsilon: float = 1e-6
  dtype: DType = jnp.float32
  scale_init: Initializer = nn.initializers.ones
  use_scale: bool = True
  conv: Optional[nn.Module] = None
  center_scale_at_zero: bool = False
  scale_axis_name: str = 'embed'

  @nn.compact
  def __call__(self,
               x: Array,
               decode: bool = False,
               prefill: bool = False,

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

# 
# Path: flaxformer/types.py



# Path: flaxformer/components/layer_norm.py
from typing import Optional
from flax import linen as nn
from flax.linen import partitioning
from jax import lax
from jax import numpy as jnp
from flaxformer.architectures.common import param_remapping
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer


class T5LayerNorm(nn.Module, param_remapping.ParameterRemappable):
  """T5 Layer normalization.

  Operates on the last axis of the input data.

  Attributes:
    epsilon: A small float added to variance to avoid dividing by zero.
    dtype: the dtype of the computation (default: float32).
    scale_init: Initializer for scale, by default, one.
    use_scale: boolean - whether to scale by a learned per-channel value
    conv: optional convolution to happen after layer norm
    center_scale_at_zero: boolean - If True, the scale weights will be
      initialized to 0.0 and a constant 1.0 will be added before they are
      applied. Preferable if you use weight decay in the optimizer.
    scale_axis_name: Axis name of the scale variable.
  """
  epsilon: float = 1e-6
  dtype: DType = jnp.float32
  scale_init: Initializer = nn.initializers.ones
  use_scale: bool = True
  conv: Optional[nn.Module] = None
  center_scale_at_zero: bool = False
  scale_axis_name: str = 'embed'

  @nn.compact
  def __call__(self,
               x: Array,
               decode: bool = False,
               prefill: bool = False,

               prefill_lengths: Optional[Array] = None) -> Array: