# 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
from prompt_tuning import prompts
from flaxformer.types import Array
import flax.linen as nn
import jax.numpy as jnp
  flat_x_embed = jnp.reshape(x_embed, (b * t, h))
  # [B, P, H] -> [B * P, H]
  flat_prompt = jnp.reshape(prompt, (b * p, h))

  # Set the data values, including padding, to the correct positions in the
  # flattened target, leaving room for the prompt.
  target = target.at[flat_data_indices, :].set(flat_x_embed)
  # Set the prompt values in the correct positions in the flattened array
  target = target.at[flat_prompt_indices, :].set(flat_prompt)
  # Reshape to the batched tensor we expect.
  # [B * (T + P), H] -> [B, (T + P), H]
  return jnp.reshape(target, (b, (t + p), h))


class Prompt(nn.Module):
  """Generate a Prompt and concatenate it with the input.

  This is the training time version of prompting a model. Calling the injected
  `prompt` module will generate your unbatched prompt. This model then
  replicates it for the batched input and concatenates them together.

  Attributes:
    prompt: The module that actually generates the unbatched prompt.
    combine: A function that combines the prompt and the embedded input.
  """
  prompt: nn.Module
  combine: Callable[[Array, Array, Array], Array] = prefix_prompt

  def __call__(self, x, x_embed):
    prompt = self.prompt(x, x_embed)

# 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/train/prompts.py
from typing import Callable
from prompt_tuning import prompts
from flaxformer.types import Array
import flax.linen as nn
import jax.numpy as jnp

  flat_x_embed = jnp.reshape(x_embed, (b * t, h))
  # [B, P, H] -> [B * P, H]
  flat_prompt = jnp.reshape(prompt, (b * p, h))

  # Set the data values, including padding, to the correct positions in the
  # flattened target, leaving room for the prompt.
  target = target.at[flat_data_indices, :].set(flat_x_embed)
  # Set the prompt values in the correct positions in the flattened array
  target = target.at[flat_prompt_indices, :].set(flat_prompt)
  # Reshape to the batched tensor we expect.
  # [B * (T + P), H] -> [B, (T + P), H]
  return jnp.reshape(target, (b, (t + p), h))


class Prompt(nn.Module):
  """Generate a Prompt and concatenate it with the input.

  This is the training time version of prompting a model. Calling the injected
  `prompt` module will generate your unbatched prompt. This model then
  replicates it for the batched input and concatenates them together.

  Attributes:
    prompt: The module that actually generates the unbatched prompt.
    combine: A function that combines the prompt and the embedded input.
  """
  prompt: nn.Module
  combine: Callable[[Array, Array, Array], Array] = prefix_prompt

  def __call__(self, x, x_embed):
    prompt = self.prompt(x, x_embed)

    prompt = prompts.expand_to_batch(prompt, x_embed)