# Path: t5x/examples/decoder_only/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 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,
#                *,
#                decode: bool = False,
#                deterministic: bool = False,
#                prefill: bool = False,
#                prefill_lengths: Optional[Array] = None) -> 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, decode=False):
#   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, Optional, Sequence
from flax import linen as nn
from flax import struct
from t5x.examples.decoder_only import layers
import jax.numpy as jnp
  dropout_rate: float = 0.1
  # If `True`, the embedding weights are used in the decoder output layer.
  logits_via_embedding: bool = False


class DecoderLayer(nn.Module):
  """Transformer decoder layer."""
  config: TransformerConfig

  @nn.compact
  def __call__(self,
               inputs: jnp.ndarray,
               decoder_mask: Optional[jnp.ndarray] = None,
               deterministic: bool = False,
               decode: bool = False,
               max_decode_length: Optional[int] = None,
               prefill: bool = False,
               prefill_lengths: Optional[jnp.ndarray] = None):
    """Applies decoder block module."""
    cfg = self.config

    # Relative position embedding as attention biases.
    l = max_decode_length if decode and max_decode_length else inputs.shape[-2]

    # During decoding, this module will be called with `decode=True` first to
    # initialize the decoder cache, including a cached relpos bias. The prefill
    # codepath will call this once again with `decode=False`, which is slightly
    # wasteful but generally harmless. During subsequent decode steps, this will
    # be called with `decode=True` and will reuse the cached bias. This
    # significantly improves performance during decoding with many decode steps.

# Path: t5x/examples/decoder_only/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 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,
#                *,
#                decode: bool = False,
#                deterministic: bool = False,
#                prefill: bool = False,
#                prefill_lengths: Optional[Array] = None) -> 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, decode=False):
#   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/decoder_only/network.py
from typing import Any, Optional, Sequence
from flax import linen as nn
from flax import struct
from t5x.examples.decoder_only import layers
import jax.numpy as jnp

  dropout_rate: float = 0.1
  # If `True`, the embedding weights are used in the decoder output layer.
  logits_via_embedding: bool = False


class DecoderLayer(nn.Module):
  """Transformer decoder layer."""
  config: TransformerConfig

  @nn.compact
  def __call__(self,
               inputs: jnp.ndarray,
               decoder_mask: Optional[jnp.ndarray] = None,
               deterministic: bool = False,
               decode: bool = False,
               max_decode_length: Optional[int] = None,
               prefill: bool = False,
               prefill_lengths: Optional[jnp.ndarray] = None):
    """Applies decoder block module."""
    cfg = self.config

    # Relative position embedding as attention biases.
    l = max_decode_length if decode and max_decode_length else inputs.shape[-2]

    # During decoding, this module will be called with `decode=True` first to
    # initialize the decoder cache, including a cached relpos bias. The prefill
    # codepath will call this once again with `decode=False`, which is slightly
    # wasteful but generally harmless. During subsequent decode steps, this will
    # be called with `decode=True` and will reuse the cached bias. This
    # significantly improves performance during decoding with many decode steps.

    decoder_bias = layers.RelativePositionBiases(