# Path: t5x/examples/t5/layers.py
# def dot_product_attention(query: Array,
#                           key: Array,
#                           value: Array,
#                           bias: Optional[Array] = None,
#                           dropout_rng: Optional[PRNGKey] = None,
#                           dropout_rate: float = 0.,
#                           deterministic: bool = False,
#                           dtype: DType = jnp.float32,
#                           float32_logits: bool = False):
#   def __call__(self,
#                inputs_q: Array,
#                inputs_kv: Array,
#                mask: Optional[Array] = None,
#                bias: Optional[Array] = None,
#                *,
#                decode: bool = False,
#                deterministic: bool = False) -> Array:
# def _normalize_axes(axes: Iterable[int], ndim: int) -> Tuple[int]:
# def _canonicalize_tuple(x):
#   def __call__(self, inputs: Array) -> Array:
# def _convert_to_activation_function(
#     fn_or_string: Union[str, Callable]) -> Callable:
#   def __call__(self, inputs, decode: bool = False, deterministic: bool = False):
#   def setup(self):
#   def __call__(self, inputs: Array) -> Array:
#   def attend(self, query: Array) -> Array:
#   def _relative_position_bucket(relative_position,
#                                 bidirectional=True,
#                                 num_buckets=32,
#                                 max_distance=128):
#   def __call__(self, qlen, klen, bidirectional=True):
#   def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
# 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:
# class MultiHeadDotProductAttention(nn.Module):
# class DenseGeneral(nn.Module):
# class MlpBlock(nn.Module):
# class Embed(nn.Module):
# class RelativePositionBiases(nn.Module):
# class LayerNorm(nn.Module):


from typing import Any, Sequence
from flax import linen as nn
from flax import struct
from t5x.examples.t5 import layers
import jax.numpy as jnp
    cfg = self.config
    assert decoder_input_tokens.ndim == 2  # [batch, len]
    rel_emb = layers.RelativePositionBiases(
        num_buckets=32,
        max_distance=128,
        num_heads=cfg.num_heads,
        dtype=cfg.dtype,
        embedding_init=nn.initializers.variance_scaling(1.0, 'fan_avg',
                                                        'uniform'),
        name='relpos_bias')

    # [batch, length] -> [batch, length, emb_dim]
    y = self.shared_embedding(decoder_input_tokens.astype('int32'))
    y = nn.Dropout(
        rate=cfg.dropout_rate, broadcast_dims=(-2,))(
            y, deterministic=deterministic)
    y = y.astype(cfg.dtype)

    for lyr in range(cfg.num_decoder_layers):
      # [batch, length, emb_dim] -> [batch, length, emb_dim]
      y = DecoderLayer(
          config=cfg, relative_embedding=rel_emb, name=f'layers_{lyr}')(
              y,
              encoded,
              decoder_mask=decoder_mask,
              encoder_decoder_mask=encoder_decoder_mask,
              deterministic=deterministic,
              decode=decode,
              max_decode_length=max_decode_length)


# Path: t5x/examples/t5/layers.py
# def dot_product_attention(query: Array,
#                           key: Array,
#                           value: Array,
#                           bias: Optional[Array] = None,
#                           dropout_rng: Optional[PRNGKey] = None,
#                           dropout_rate: float = 0.,
#                           deterministic: bool = False,
#                           dtype: DType = jnp.float32,
#                           float32_logits: bool = False):
#   def __call__(self,
#                inputs_q: Array,
#                inputs_kv: Array,
#                mask: Optional[Array] = None,
#                bias: Optional[Array] = None,
#                *,
#                decode: bool = False,
#                deterministic: bool = False) -> Array:
# def _normalize_axes(axes: Iterable[int], ndim: int) -> Tuple[int]:
# def _canonicalize_tuple(x):
#   def __call__(self, inputs: Array) -> Array:
# def _convert_to_activation_function(
#     fn_or_string: Union[str, Callable]) -> Callable:
#   def __call__(self, inputs, decode: bool = False, deterministic: bool = False):
#   def setup(self):
#   def __call__(self, inputs: Array) -> Array:
#   def attend(self, query: Array) -> Array:
#   def _relative_position_bucket(relative_position,
#                                 bidirectional=True,
#                                 num_buckets=32,
#                                 max_distance=128):
#   def __call__(self, qlen, klen, bidirectional=True):
#   def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
# 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:
# class MultiHeadDotProductAttention(nn.Module):
# class DenseGeneral(nn.Module):
# class MlpBlock(nn.Module):
# class Embed(nn.Module):
# class RelativePositionBiases(nn.Module):
# class LayerNorm(nn.Module):


# Path: t5x/examples/t5/network.py
from typing import Any, Sequence
from flax import linen as nn
from flax import struct
from t5x.examples.t5 import layers
import jax.numpy as jnp

    cfg = self.config
    assert decoder_input_tokens.ndim == 2  # [batch, len]
    rel_emb = layers.RelativePositionBiases(
        num_buckets=32,
        max_distance=128,
        num_heads=cfg.num_heads,
        dtype=cfg.dtype,
        embedding_init=nn.initializers.variance_scaling(1.0, 'fan_avg',
                                                        'uniform'),
        name='relpos_bias')

    # [batch, length] -> [batch, length, emb_dim]
    y = self.shared_embedding(decoder_input_tokens.astype('int32'))
    y = nn.Dropout(
        rate=cfg.dropout_rate, broadcast_dims=(-2,))(
            y, deterministic=deterministic)
    y = y.astype(cfg.dtype)

    for lyr in range(cfg.num_decoder_layers):
      # [batch, length, emb_dim] -> [batch, length, emb_dim]
      y = DecoderLayer(
          config=cfg, relative_embedding=rel_emb, name=f'layers_{lyr}')(
              y,
              encoded,
              decoder_mask=decoder_mask,
              encoder_decoder_mask=encoder_decoder_mask,
              deterministic=deterministic,
              decode=decode,
              max_decode_length=max_decode_length)


    y = layers.LayerNorm(dtype=cfg.dtype, name='decoder_norm')(y)