# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
#   """Expand unbatched `x` to the same batch size as `y`."""
#   batch_size = y.shape[0]
#   return jnp.tile(
#       jnp.expand_dims(x, axis=0), [batch_size] + [1 for _ in x.shape])
# 
# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
# def np_load(path: str, restore_dtype: Optional[DType] = None) -> Array:
# def t5x_load(checkpoint_path: str,
#              variable_path: str,
#              restore_dtype: Optional[DType] = None) -> Array:
# def from_array(prompt: Array) -> Initializer:
#   def initialize_from_file(unused_prng_key: Array,
#                            shape: Sequence[int]) -> Array:
# def from_sample_of_embeddings(
#     embeddings: Array,
#     population_size: Optional[int] = None,
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedding_sample(rng: Array,
#                                        shape: Sequence[int]) -> Array:
# def from_embedded_list(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     texts: Sequence[str],
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_list(rng: Array, shape: Sequence[int]) -> Array:
# def from_embedded_string(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     text: str,
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_string(rng: Array,
#                                       shape: Sequence[int]) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


from typing import Callable, Optional, Any, Iterable, Tuple
from flax import linen as nn
from flax.linen import partitioning
from prompt_tuning.prompts import expand_to_batch
from prompt_tuning.prompts import Initializer
from flaxformer.components import embedding
from flaxformer.types import Array
import jax.numpy as jnp
  `get_task_size` that are overridden in subclasses. We were not able to get
  this to work. When running inside of T5X on TPU, we couldn't stop JAX from
  outputting traced values for methods like `get_shared_features`, despite it
  being based on shape information only. This caused errors when we used the
  result in the prompt shape.

  TODO: Create a minimum reproducible example in colab, file a bug,
  and link that bug here.

  We solved this by turning my methods into explicit functions that are class
  attributes (with defaults). When these functions are called, the module's
  `self` is explicitly passed in. This replicates having methods on subclasses
  but avoids the JAX error and allows us to use the results of computation for
  the parameter sizes.
"""



GetFeatures = Optional[Callable[[nn.Module, int], int]]
GetTaskSize = Optional[Callable[[nn.Module, int, int], int]]
Combine = Optional[Callable[[Array, Array], Array]]


def identity(x: Any) -> Any:
  """Return `x`, used as default function."""
  return x


def sentinel_to_task(x: Array, vocab_size: int) -> Array:
  r"""Convert the vocabulary index of tokens like <extra_id\d+> to 0 offsets."""

# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
#   """Expand unbatched `x` to the same batch size as `y`."""
#   batch_size = y.shape[0]
#   return jnp.tile(
#       jnp.expand_dims(x, axis=0), [batch_size] + [1 for _ in x.shape])
# 
# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
# def np_load(path: str, restore_dtype: Optional[DType] = None) -> Array:
# def t5x_load(checkpoint_path: str,
#              variable_path: str,
#              restore_dtype: Optional[DType] = None) -> Array:
# def from_array(prompt: Array) -> Initializer:
#   def initialize_from_file(unused_prng_key: Array,
#                            shape: Sequence[int]) -> Array:
# def from_sample_of_embeddings(
#     embeddings: Array,
#     population_size: Optional[int] = None,
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedding_sample(rng: Array,
#                                        shape: Sequence[int]) -> Array:
# def from_embedded_list(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     texts: Sequence[str],
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_list(rng: Array, shape: Sequence[int]) -> Array:
# def from_embedded_string(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     text: str,
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_string(rng: Array,
#                                       shape: Sequence[int]) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


# Path: prompt_tuning/extended/multitask_prompts.py
from typing import Callable, Optional, Any, Iterable, Tuple
from flax import linen as nn
from flax.linen import partitioning
from prompt_tuning.prompts import expand_to_batch
from prompt_tuning.prompts import Initializer
from flaxformer.components import embedding
from flaxformer.types import Array
import jax.numpy as jnp

  `get_task_size` that are overridden in subclasses. We were not able to get
  this to work. When running inside of T5X on TPU, we couldn't stop JAX from
  outputting traced values for methods like `get_shared_features`, despite it
  being based on shape information only. This caused errors when we used the
  result in the prompt shape.

  TODO: Create a minimum reproducible example in colab, file a bug,
  and link that bug here.

  We solved this by turning my methods into explicit functions that are class
  attributes (with defaults). When these functions are called, the module's
  `self` is explicitly passed in. This replicates having methods on subclasses
  but avoids the JAX error and allows us to use the results of computation for
  the parameter sizes.
"""



GetFeatures = Optional[Callable[[nn.Module, int], int]]
GetTaskSize = Optional[Callable[[nn.Module, int, int], int]]
Combine = Optional[Callable[[Array, Array], Array]]


def identity(x: Any) -> Any:
  """Return `x`, used as default function."""
  return x


def sentinel_to_task(x: Array, vocab_size: int) -> Array:
  r"""Convert the vocabulary index of tokens like <extra_id\d+> to 0 offsets."""

  return vocab_size - 1 - x