# 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

  Args:
    shared: [P_s, H] A prompt that is shared across tasks (unbatched).
    tasks: [B, P_t * H] A batch of task specific prompts.

  Returns:
    Task specific prompts concatenated to a shared prompt along the sequence
    dimension. [B, P_s + P_t, H]
  """
  _, size = shared.shape
  task_length = tasks.shape[1] // size
  tasks = jnp.reshape(tasks, (-1, task_length, size))
  shared = expand_to_batch(shared, tasks)
  return jnp.concatenate([shared, tasks], axis=1)


def concat_features_shared_and_tasks(shared: Array, tasks: Array) -> Array:
  """Combine shared and task prompts via feature vector concatenation.

  Args:
    shared: [P, H_s] A prompt that is shared across tasks (unbatched).
    tasks: [B, P * H_t] A batch of task specific prompts.

  Returns:
    Task specific prompts concatenated to a shared prompts along the feature
    dimension. [B, P, H_s + H_t] s.t. H_s + H_t = H
  """
  length, _ = shared.shape
  task_size = tasks.shape[1] // length
  tasks = jnp.reshape(tasks, (-1, length, task_size))

# 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


  Args:
    shared: [P_s, H] A prompt that is shared across tasks (unbatched).
    tasks: [B, P_t * H] A batch of task specific prompts.

  Returns:
    Task specific prompts concatenated to a shared prompt along the sequence
    dimension. [B, P_s + P_t, H]
  """
  _, size = shared.shape
  task_length = tasks.shape[1] // size
  tasks = jnp.reshape(tasks, (-1, task_length, size))
  shared = expand_to_batch(shared, tasks)
  return jnp.concatenate([shared, tasks], axis=1)


def concat_features_shared_and_tasks(shared: Array, tasks: Array) -> Array:
  """Combine shared and task prompts via feature vector concatenation.

  Args:
    shared: [P, H_s] A prompt that is shared across tasks (unbatched).
    tasks: [B, P * H_t] A batch of task specific prompts.

  Returns:
    Task specific prompts concatenated to a shared prompts along the feature
    dimension. [B, P, H_s + H_t] s.t. H_s + H_t = H
  """
  length, _ = shared.shape
  task_size = tasks.shape[1] // length
  tasks = jnp.reshape(tasks, (-1, length, task_size))

  shared = expand_to_batch(shared, tasks)