# 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
# def prefix_prompt(prompt: Array, x_embed: Array, x: Array) -> Array:
# def prefix_prompt_after_bos(prompt: Array, x_embed: Array, x: Array) -> Array:
# def suffix_prompt(prompt: Array,
#                   x_embed: Array,
#                   x: Array,
#                   decoder_only: bool = False,
#                   before_eos: bool = True) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning import prompts
from prompt_tuning.train import prompts as train_prompts
import jax
import jax.numpy as jnp
import numpy as np
      inputs[:, 0] = 0
    inputs = jnp.asarray(inputs)
    embed = jnp.asarray(embed)
    with_prompt = jax.jit(train_prompts.suffix_prompt, static_argnums=(3, 4))(
        prompt, embed, inputs, decoder_only=decoder_only, before_eos=before_eos)
    self.assertEqual(with_prompt.shape,
                     (batch_size, seq_len + prompt_length, embed_size))
    for i, (example, l) in enumerate(zip(with_prompt, lengths)):
      if before_eos:
        np.testing.assert_array_equal(example[:l - 1], embed[i, :l - 1])
        np.testing.assert_array_equal(example[l + prompt_length], embed[i, l])
        np.testing.assert_array_equal(example[l - 1:l - 1 + prompt_length],
                                      prompt[i])
      else:
        np.testing.assert_array_equal(example[:l], embed[i, :l])
        np.testing.assert_array_equal(example[l:l + prompt_length], prompt[i])
      padding = example[l + prompt_length:]
      np.testing.assert_array_equal(padding,
                                    np.full_like(padding, padding_embed_value))

  def test_prompt_does_concatenation(self):
    embed_size = 20
    prompt_length = 5
    batch_size = 2
    seq_len = 20
    mock_prompt = mock.create_autospec(
        prompts.Prompt, spec_set=True, instance=True)
    prompt = jnp.zeros((prompt_length, embed_size))
    mock_prompt.return_value = prompt
    mock_combine = mock.create_autospec(

# 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
# def prefix_prompt(prompt: Array, x_embed: Array, x: Array) -> Array:
# def prefix_prompt_after_bos(prompt: Array, x_embed: Array, x: Array) -> Array:
# def suffix_prompt(prompt: Array,
#                   x_embed: Array,
#                   x: Array,
#                   decoder_only: bool = False,
#                   before_eos: bool = True) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


# Path: prompt_tuning/train/prompts_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning import prompts
from prompt_tuning.train import prompts as train_prompts
import jax
import jax.numpy as jnp
import numpy as np

      inputs[:, 0] = 0
    inputs = jnp.asarray(inputs)
    embed = jnp.asarray(embed)
    with_prompt = jax.jit(train_prompts.suffix_prompt, static_argnums=(3, 4))(
        prompt, embed, inputs, decoder_only=decoder_only, before_eos=before_eos)
    self.assertEqual(with_prompt.shape,
                     (batch_size, seq_len + prompt_length, embed_size))
    for i, (example, l) in enumerate(zip(with_prompt, lengths)):
      if before_eos:
        np.testing.assert_array_equal(example[:l - 1], embed[i, :l - 1])
        np.testing.assert_array_equal(example[l + prompt_length], embed[i, l])
        np.testing.assert_array_equal(example[l - 1:l - 1 + prompt_length],
                                      prompt[i])
      else:
        np.testing.assert_array_equal(example[:l], embed[i, :l])
        np.testing.assert_array_equal(example[l:l + prompt_length], prompt[i])
      padding = example[l + prompt_length:]
      np.testing.assert_array_equal(padding,
                                    np.full_like(padding, padding_embed_value))

  def test_prompt_does_concatenation(self):
    embed_size = 20
    prompt_length = 5
    batch_size = 2
    seq_len = 20
    mock_prompt = mock.create_autospec(
        prompts.Prompt, spec_set=True, instance=True)
    prompt = jnp.zeros((prompt_length, embed_size))
    mock_prompt.return_value = prompt
    mock_combine = mock.create_autospec(

        train_prompts.prefix_prompt, spec_set=True)