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: _DEFAULT_LAYER_NORM = 1e-12
_DEFAULT_INIT_RANGE = 0.02
class BertPooler(nn.Module):
class ClassifierHead(nn.Module):
class NSPHead(nn.Module):
class MLMHead(nn.Module):
  def _extract_cls_embedding(self, encoded_inputs: Array) -> Array:
  def __call__(self, encoded_inputs: Array, **unused_kwargs):
  def setup(self):
  def __call__(self,
               encoded_inputs: Array,
               *,
               enable_dropout: Optional[bool] = None,
               **unused_kwargs) -> Array:
  def setup(self):
  def __call__(self, encoded_inputs: Array, **unused_kwargs) -> Array:
def gather_indices(inputs: Array, indices: Array) -> Array:
  def setup(self):
  def __call__(self, encoded_inputs: Array, *,
               masked_positions: Optional[Array]) -> Array: def dense(features, name, inputs, kernel_axis_names):
  if self.weight_params is not None or self.act_params is not None:
    # TODO: Push the "quantized vs not" decision down into the
    # AQT library. Currently we make that decision here, because the AQT
    # library doesn't support DenseGeneral, so there's extra reshapes here
    # whose performance impact I don't know.
    aqt_context = aqt_config.QuantContext(
        update_bounds=False, collect_acts_stats=False)
    aqt_hparams = aqt_flax_layers.DenseAqt.HParams(
        weight_prec=self.weight_params.prec,
        weight_half_shift=self.weight_params.half_shift,
        quant_act=self.act_params,  # currently supports fixed bounds only.
        quant_type=aqt.QuantType.aqt,
        weight_quant_granularity=aqt_config.QuantGranularity.per_channel,
    )
    batch, seq_len, channels = inputs.shape
    inputs = inputs.reshape((batch * seq_len, channels))
    result = aqt_flax_layers.DenseAqt(
        features=features,
        hparams=aqt_hparams,
        train=enable_dropout,
        quant_context=aqt_context,
        paxis_name=None,
        # No "cross-replica" reduction expressed in the XLA graph at this
        # stage. Will be imposed later, automatically, by XLA SPMD.
        use_bias=self.use_bias,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        dtype=self.dtype,
        kernel_axis_names=kernel_axis_names,
        name=name,
    )(inputs, padding_mask=None)
    return result.reshape((batch, seq_len, features))
  else:
    return DenseGeneral(
        features,
        use_bias=self.use_bias,
        dtype=self.dtype,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        kernel_axis_names=kernel_axis_names,
        name=name)(
            inputs) 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: def sinusoidal(min_scale: float = 1.0,
               max_scale: float = 10000.0,
               dtype: DType = jnp.float32) -> Initializer:
  def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
def truncated_normal(mean: Union[float, Array] = 0.0,
                     stddev: Union[float, Array] = 0.05,
                     dtype: DType = jnp.float32) -> Initializer:
  def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array: 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):
from typing import Optional, Tuple
from flax import linen as nn
from flaxformer import transformer_common as common
from flaxformer.architectures.bert import heads
from flaxformer.components import dense
from flaxformer.components import embedding
from flaxformer.components import initializers
from flaxformer.components.attention import dense_attention
from flaxformer.types import Activation
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
import chex
import jax.numpy as jnp
  encoder_block: EncoderBlock

  def __call__(self,
               input_ids: Array,
               *,
               input_mask: Array,
               position_ids: Optional[Array] = None,
               segment_ids: Optional[Array] = None,
               enable_dropout: bool = True) -> Array:
    """Embeds the inputs and then encodes those representations.

    Args:
      input_ids: <int>[batch..., seq_len]. The first position in the sequence
        should correspond to a "beginning of input" symbol, generally `[CLS]`.
      input_mask: <int>[batch..., seq_len]. Indicates which positions in
        `input_ids` are non-padding (0 for padding, 1 otherwise).
      position_ids: <int>[batch..., seq_len]. The position of each input ID
        within its sequence. This is typically just `range(seq_len)` for each
        sequence, but custom behavior may be desired if, for instance, multiple
        examples are packed into each sequence.
      segment_ids: <int>[batch..., seq_len]. Indicates the "type" of each input
        position. For a traditional BERT-style model with two segments, valid
        values would be {0, 1}.
      enable_dropout: Enables dropout when set to True.

    Returns:
      <float>[batch..., seq_len, hidden_size].
    """
    # Validate inputs and create variables for dimension sizes.
    chex.assert_shape(input_ids, (..., None))
    chex.assert_type(input_ids, int)
    *batch_sizes, seq_len = input_ids.shape
    chex.assert_shape(input_mask, (*batch_sizes, seq_len))
    if position_ids is not None:
      chex.assert_shape(position_ids, (*batch_sizes, seq_len))
      chex.assert_type(position_ids, int)
    if segment_ids is not None:
      chex.assert_shape(segment_ids, (*batch_sizes, seq_len))
      chex.assert_type(segment_ids, int)

    embeddings = self.embedder_block(
        input_ids=input_ids,
        position_ids=position_ids,
        segment_ids=segment_ids,
        enable_dropout=enable_dropout)
    chex.assert_shape(embeddings, (*batch_sizes, seq_len, None))

    attention_mask = dense_attention.make_attention_mask(input_mask, input_mask)
    chex.assert_shape(attention_mask, (*batch_sizes, 1, seq_len, seq_len))

    result = self.encoder_block(
        embeddings,
        attention_mask=attention_mask,
        enable_dropout=enable_dropout)
    chex.assert_shape(result, (*batch_sizes, seq_len, None))

    return result


class EmbedderBlock(nn.Module):
  """Embeds the inputs, then applies dropout and layer norm."""

  embedder: embedding.MultiEmbed
  dropout: nn.Dropout
  layer_norm: nn.LayerNorm

  def __call__(self,
               input_ids: Array,
               *,
               position_ids: Optional[Array] = None,
               segment_ids: Optional[Array] = None,
               enable_dropout: bool = True) -> Array:
    """Embeds the inputs, then applies dropout and layer norm.

    Args:
      input_ids: <int>[batch..., seq_len]. The first position in the sequence
        should correspond to a "beginning of input" symbol, generally `[CLS]`.
      position_ids: <int>[batch..., seq_len]. The position of each input ID
        within its sequence. This is typically just `range(seq_len)` for each
        sequence, but custom behavior may be desired if, for instance, multiple
        examples are packed into each sequence.
      segment_ids: <int>[batch..., seq_len]. Indicates the "type" of each input
        position. For a traditional BERT-style model with two segments, valid
        values would be {0, 1}.
      enable_dropout: Enables dropout when set to True.

    Returns:
      <float>[batch..., seq_len, hidden_size].
    """
    # Validate inputs and create variables for dimension sizes.
    chex.assert_shape(input_ids, (..., None))
    chex.assert_type(input_ids, int)
    *batch_sizes, seq_len = input_ids.shape
    if position_ids is not None:
      chex.assert_shape(position_ids, (*batch_sizes, seq_len))
      chex.assert_type(position_ids, int)
    if segment_ids is not None:
      chex.assert_shape(segment_ids, (*batch_sizes, seq_len))
      chex.assert_type(segment_ids, int)

    embeddings = self.embedder(
        input_ids=input_ids,
        position_ids=position_ids,
        segment_ids=segment_ids,
    )
    chex.assert_shape(embeddings, (*batch_sizes, seq_len, None))

    embeddings = self.layer_norm(embeddings)
    embeddings = self.dropout(embeddings, deterministic=not enable_dropout)
    chex.assert_shape(embeddings, (*batch_sizes, seq_len, None))

    return embeddings


class EncoderBlock(nn.Module):
  """A BERT encoder block: a sequence of transformer layers.

  Passes inputs through the full stack of encoder layers. Note that dropout and
  layer norm are performed within and at the end of each encoder layer.
  """

  layer_sequence: common.LayerSequence