def global_mesh_defined():
def with_sharding(x, partitioning_dims: int):
def with_sharding_migration(
    x: T,
    activation_partitioning_dims: Optional[int],
    logical_axis_names: Tuple[str, ...],
) -> T:
T = TypeVar('T') class LayerSequence(nn.Module, param_remapping.ParameterRemappable):
class TransparentLayerSequence:
  def setup(self):
  def __call__(self, inputs: Array, *args, **kwargs) -> Array:
  def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],
                            inputs: Array, *args, **kwargs) -> Array:
  def __call__(self, inputs: Array, *args, **kwargs) -> Array:
  def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],
                            inputs: Array, *args, **kwargs) -> Array: 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]: class EmbedCombineMethod(enum.IntEnum):
class Embedder(Generic[_Inputs], metaclass=abc.ABCMeta):
class InspectableMultiEmbedder(Generic[_Inputs], Embedder[_Inputs]):
class DictEmbedder(nn.Module, InspectableMultiEmbedder[Mapping[str, Any]]):
class EmbedderWithDecode(Generic[_Inputs], Embedder[_Inputs]):
class EmbedderWithDeterministic(Generic[_Inputs], Embedder[_Inputs]):
class Embed(nn.Module, Embedder[Array]):
class MultiEmbed(nn.Module):
class FixedEmbed(nn.Module, EmbedderWithDecode[Array]):
class PositionEmbed(nn.Module, EmbedderWithDecode[Array]):
class HashEmbed(nn.Module, Embedder[Array]):
class NgramHashEmbed(nn.Module, Embedder[Array]):
  SUM = 1
  CONCAT = 2
  def __call__(self,
               inputs: _Inputs,
               *,
               segment_ids: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True) -> Array:
  def __call__(self,
               inputs: _Inputs,
               *,
               segment_ids: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True) -> Array:
  def get_individual_embeddings(self,
                                inputs: Array,
                                *,
                                segment_ids: Optional[Array] = None,
                                decode: bool = False,
                                enable_dropout: bool = True) -> Sequence[Array]:
  def combine_embeddings(self, embeddings: Sequence[Array]) -> Array:
  def get_individual_embeddings(self,
                                inputs: Mapping[str, Any],
                                *,
                                segment_ids: Optional[Array] = None,
                                decode: bool = False,
                                enable_dropout: bool = True) -> Sequence[Array]:
  def combine_embeddings(self, embeddings: Sequence[Array]) -> Array:
  def setup(self):
  def __call__(
      self,
      inputs: Array,
      *,
      segment_ids: Optional[Array] = None,
      decode: bool = False,
      enable_dropout: bool = True,
      input_axis_names: Optional[Sequence[str]] = ('batch', 'length'),
  ) -> Array:
  def attend(self, query: Array) -> Array:
  def get_individual_embeddings(
      self,
      decode: bool = False,
      deterministic: bool = False,
      segment_ids: Optional[Array] = None,
      **input_kwargs: Mapping[str, Array]) -> Dict[str, Array]:
  def combine_embeddings(
      self,
      embeddings: Dict[str, Array],
      combine_method: EmbedCombineMethod = EmbedCombineMethod.SUM) -> Array:
  def __call__(self,
               combine_method: EmbedCombineMethod = EmbedCombineMethod.SUM,
               *,
               segment_ids: Optional[Array] = None,
               decode: bool = False,
               deterministic: bool = False,
               **input_kwargs: Mapping[str, Array]) -> Array:
  def setup(self):
  def __call__(self,
               inputs,
               *,
               segment_ids: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True):
  def setup(self):
  def __call__(self,
               inputs,
               *,
               segment_ids: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True):
def rotate_half(x):
def apply_rotary_embedding(q,
                           k,
                           cos,
                           sin,
                           batch_size=None,
                           num_heads=None,
                           decode=False,
                           rotary_index=None):
def generate_fixed_pos_embedding(features,
                                 length,
                                 min_timescale=1.0,
                                 max_timescale=10000.0):
def _generate_primes(num_primes: int) -> Sequence[int]:
def _is_prime(n: int) -> int:
  def setup(self):
  def __call__(self,
               input_ids: Array,
               *,
               segment_ids: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True) -> Array:
  def setup(self):
  def __call__(self,
               input_ids: Array,
               *,
               segment_ids: Optional[Array],
               decode: bool = False,
               enable_dropout: bool = True) -> Array:
  def _shift_left(self, ids: Array) -> Array: class RelativePositionBiases(nn.Module):
  def _relative_position_bucket(relative_position,
                                bidirectional=True,
                                num_buckets=32,
                                max_distance=128):
  def __call__(self, qlen, klen, bidirectional=True, decode=False): class RichAttentionApi(metaclass=abc.ABCMeta):
  def __call__(self,
               q_inputs: Array,
               k_inputs: Array,
               bidirectional: bool = True,
               is_cross_attention: bool = False) -> Array: def core_remat_static(fn,
                      variables=True,
                      rngs=True,
                      concrete=False,
                      prevent_cse=True,
                      static_argnums=(),
                      policy=None):
  def _repack_remat_args(dyn_args, static_args):
  def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):
    def rematted(variable_groups, rng_groups, *dyn_args):
def remat(target,
          variables=True,
          rngs=True,
          concrete=False,
          prevent_cse=True,
          static_argnums=(),
          policy=None,
          methods=None):
def canonicalize_arguments(orig_fn):
  def dekwarged_fn(*args):
  def convert_to_args(*args, **kwargs):
def canonicalized_class_transform(trafo, clz, *t_args, **t_kwargs):
  def post_fn(self, *args, **kwargs):
def apply_transform_to_module_factory(trafo, factory, *args, **kwargs):
  def new_factory():
def global_mesh_defined():
def with_sharding_constraint(x, axis_resources):
def inner_scan_spmd(annotation_tree, scan_axis):
  def del_axis(x):
  def annotate_fn(variable_groups, rng_groups):
    def maybe_annotate_group(x): def _softmax_with_extra_logit(
    x: Array,
    axis: Optional[Union[int, Tuple[int, ...]]] = -1,
) -> Array:
def dot_product_attention_weights(query: Array,
                                  key: Array,
                                  bias: Optional[Array] = None,
                                  broadcast_dropout: bool = True,
                                  rescale_logits: bool = False,
                                  dropout_rng: Optional[PRNGKey] = None,
                                  dropout_rate: float = 0.,
                                  enable_dropout: bool = True,
                                  dtype: DType = jnp.float32,
                                  precision: Optional[lax.Precision] = None,
                                  use_extra_logit: bool = False,
                                  float32_logits: bool = False) -> Array:
def apply_dot_product_attention_weights_to_values(
    attention_weights: Array,
    value: Array,
    precision: Optional[lax.Precision] = None) -> Array:
def dot_product_attention(query: Array,
                          key: Array,
                          value: Array,
                          bias: Optional[Array] = None,
                          broadcast_dropout: bool = True,
                          rescale_logits: bool = False,
                          dropout_rng: Optional[PRNGKey] = None,
                          dropout_rate: float = 0.,
                          enable_dropout: bool = True,
                          dtype: DType = jnp.float32,
                          precision: Optional[lax.Precision] = None,
                          use_extra_logit: bool = False,
                          float32_logits: bool = False):
def dot_product_attention_multiquery(query: Array,
                                     key: Array,
                                     value: Array,
                                     bias: Optional[Array] = None,
                                     broadcast_dropout: bool = True,
                                     rescale_logits: bool = False,
                                     dropout_rng: Optional[PRNGKey] = None,
                                     dropout_rate: float = 0.,
                                     enable_dropout: bool = True,
                                     dtype: DType = jnp.float32,
                                     precision: Optional[lax.Precision] = None,
                                     use_extra_logit: bool = False,
                                     float32_logits: bool = False) -> Array:
  def __call__(self,
               inputs_q: Array,
               inputs_kv: Array,
               mask: Optional[Array] = None,
               bias: Optional[Array] = None,
               *,
               precomputed_qkv: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True) -> Array:
  def update_cache_prefill(
      self, key: Array, value: Array, cached_key: variables.Variable,
      cached_value: variables.Variable, cache_index: variables.Variable,
      prefill_lengths: Array
  ) -> Tuple[Array, Array, Array, Array, Array, Array]:
  def update_cache_decode(
      self, key: Array, value: Array, cached_key: variables.Variable,
      cached_value: variables.Variable, cache_index: variables.Variable
  ) -> Tuple[Array, Array, Array, Array, Array, Array]:
  def __call__(self,
               inputs_q: Array,
               inputs_kv: Array,
               mask: Optional[Array] = None,
               bias: Optional[Array] = None,
               *,
               precomputed_qkv: Optional[Array] = None,
               decode: bool = False,
               enable_dropout: bool = True,
               prefill: bool = False,
               prefill_lengths: Optional[Array] = None) -> Array:
  def update_cache_prefill(
      self, key: Array, value: Array, cached_key: variables.Variable,
      cached_value: variables.Variable, cache_index: variables.Variable,
      prefill_lengths: Array
  ) -> Tuple[Array, Array, Array, variables.Variable, variables.Variable,
  def update_cache_decode(
      self, key: Array, value: Array, cached_key: variables.Variable,
      cached_value: variables.Variable, cache_index: variables.Variable
  ) -> Tuple[Array, Array, Array, variables.Variable, variables.Variable,
  def __call__(self,
               inputs_q: Array,
               inputs_kv: Array,
               mask: Optional[Array] = None,
               bias: Optional[Array] = None,
               *,
               precomputed_qkv=None,
               decode: bool = False,
               enable_dropout: bool = True,
               prefill: bool = False,
               prefill_lengths: Optional[Array] = None) -> Array:
  def setup(self) -> None:
  def __call__(
      self,
      inputs_q: Array,
      inputs_kv: Array,
      mask: Optional[Array] = None,
      bias: Optional[Array] = None,
      *,
      decode: bool = False,
      enable_dropout: bool = True,
  ) -> Array:
def make_attention_mask(query_input: Array,
                        key_input: Array,
                        pairwise_fn: Callable = jnp.multiply,
                        extra_batch_dims: int = 0,
                        dtype: DType = jnp.float32) -> Array:
def make_causal_mask(x: Array,
                     extra_batch_dims: int = 0,
                     dtype: DType = jnp.float32) -> Array:
def combine_masks(*masks: Optional[Array], dtype: DType = jnp.float32):
def combine_biases(*masks: Optional[Array]):
def make_decoder_mask(decoder_target_tokens: Array,
                      dtype: DType,
                      decoder_causal_attention: Optional[Array] = None,
                      decoder_segment_ids: Optional[Array] = None) -> Array:
def validate_dense_attention_call_parameter_shapes(inputs_q: Array,
                                                   inputs_kv: Array,
                                                   mask: Optional[Array],
                                                   bias: Optional[Array],
                                                   num_heads: Optional[int]):
def shift_left(x, axis=-1):
def get_decoder_logit_mask(decoder_input_tokens, dtype):
class DenseAttention(metaclass=abc.ABCMeta):
class MultiHeadDotProductAttention(nn.Module, DenseAttention):
class MultiQueryDotProductAttention(nn.Module, DenseAttention):
class LocalAttentionLayer(nn.Module, DenseAttention):
import inspect
import jax
import jax.numpy as jnp
from typing import Callable, Optional, Any, Tuple
from flax import linen as nn
from typing_extensions import Protocol
from flaxformer import activation_partitioning
from flaxformer import transformer_common as common
from flaxformer.architectures.common import param_remapping
from flaxformer.components import embedding
from flaxformer.components import relative_position_biases
from flaxformer.components import rich_attention_position_scores
from flaxformer.components import transforms
from flaxformer.components.attention import dense_attention
from flaxformer.types import Array
from flaxformer.types import DType
          self.mlp(x, enable_dropout=enable_dropout))
      y *= 2**-0.5
      y = layer_input + self.dropout(y, deterministic=not enable_dropout)

    else:
      # Attention block.
      x = self.pre_attention_layer_norm(layer_input)
      x = activation_partitioning.with_sharding_migration(
          x,
          self.activation_partitioning_dims,
          logical_axis_names=('batch', 'length', 'embed'))

      if logit_mask is not None:
        x = logit_mask * x

      encoder_bias = self.get_bias(layer_input=x)

      # The shape should be maintained for the residual connection.
      # [batch, length, emb_dim] -> [batch, length, emb_dim]
      x = self.attention(
          x, x, encoder_mask, encoder_bias, enable_dropout=enable_dropout)
      x = layer_input + self.post_attention_dropout(
          x, deterministic=not enable_dropout)
      x = activation_partitioning.with_sharding_migration(
          x,
          self.activation_partitioning_dims,
          logical_axis_names=('batch', 'length', 'embed'))

      # MLP block.
      y = self.pre_mlp_layer_norm(x)
      y = activation_partitioning.with_sharding_migration(
          y,
          self.activation_partitioning_dims,
          logical_axis_names=('batch', 'length', 'embed'))

      if logit_mask is not None:
        y = logit_mask * y

      # [batch, length, emb_dim] -> [batch, length, emb_dim]
      y = self.mlp(y, enable_dropout=enable_dropout)
      y = x + self.post_mlp_dropout(y, deterministic=not enable_dropout)
    y = activation_partitioning.with_sharding_migration(
        y,
        self.activation_partitioning_dims,
        logical_axis_names=('batch', 'length', 'embed'))
    if self.sow_intermediates:
      self.sow('intermediates', 'activations', y)

    # scan expects functions to have a signature: fn(carry, in) --> carry, out
    # TODO: automate this detail.
    if self.scanned:
      return y, None
    else:
      return y


class DecoderLayer(nn.Module, param_remapping.ParameterRemappable):
  """Transformer encoder-decoder layer.

  Attributes:
    self_attention: An instance of a self-attention module.
    encoder_decoder_attention: Encoder-decoder attention module. This must be
      non-None if attending to encoded representations.
    mlp: The MLP module, applied after both attention modules.
    dropout_factory:  A callable that returns a new dropout instance. This is
      applied after the attention module.
    layer_norm_factory:  A callable that returns a new layer norm. This is
      applied before the attention module and before the MLP.
    relative_position_bias_factory: A callable that returns relative position
      bias instances. This should only be used for per-layer relative position
      biases; please use `shared_relative_position_bias` if they are shared
      among layers.
    shared_relative_position_bias: An instance of a shared relative position
      bias module, usually owned by the Decoder.
    activation_partitioning_dims: When set to 2, partitions intermediate
      variables containing the input and output of the decoder layer.
    parallel: whether to call attention and mlp in parallel
    sow_intermediates: whether to track intermediates using Module.sow.
  """
  self_attention: nn.Module
  encoder_decoder_attention: Optional[nn.Module]
  mlp: nn.Module
  dropout_factory: Callable[[], nn.Module]
  layer_norm_factory: Callable[[], nn.Module]
  relative_position_bias_factory: Optional[Callable[[], nn.Module]] = None
  shared_relative_position_bias: Optional[nn.Module] = None
  activation_partitioning_dims: int = 1
  parallel: bool = False
  sow_intermediates: bool = False
  scanned: bool = False

  def setup(self):
    if (self.relative_position_bias_factory is not None and
        self.shared_relative_position_bias is not None):
      raise ValueError(
          'Please set at most one of relative_position_bias_factory and shared_relative_position_bias. '
          '(They can both be None however, e.g. for absolute position embeds.)')
    self.relpos_bias = (
        self.relative_position_bias_factory()
        if self.relative_position_bias_factory is not None else
        self.shared_relative_position_bias)

    if self.parallel:
      self.layer_norm = self.layer_norm_factory()
      self.dropout = self.dropout_factory()
    else:
      self.pre_self_attention_layer_norm = self.layer_norm_factory()
      self.post_self_attention_dropout = self.dropout_factory()
      self.pre_cross_attention_layer_norm = self.layer_norm_factory()
      self.post_cross_attention_dropout = self.dropout_factory()
      self.pre_mlp_layer_norm = self.layer_norm_factory()
      self.post_mlp_dropout = self.dropout_factory()

  def get_bias(self, max_decode_length: Optional[int], decode: bool,
               layer_input: Array,
               encoded: Array) -> Tuple[Optional[Array], Optional[Array]]:
    decoder_bias = None
    encoder_decoder_bias = None
    if self.relpos_bias:
      if isinstance(self.relpos_bias,

                    relative_position_biases.RelativePositionBiases):