# 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):


import os
import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import seqio
from typing import List
from unittest import mock
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning import prompts
from tensorflow.io import gfile
      choice_mock.assert_called_once_with(rng, population_size,
                                          (prompt_length,), False)

  @parameterized.named_parameters([
      dict(testcase_name="None", population_size=None),
      dict(testcase_name="too_big", population_size=1000)
  ])
  def test_from_sample_of_embeddings_population_size_mismatch(
      self, population_size):
    vocab_size = 100
    embed_size = 10
    prompt_length = 15
    embeddings = np.zeros((vocab_size, embed_size))

    init = prompts.from_sample_of_embeddings(embeddings, population_size)
    with mock.patch.object(
        prompts.jax.random, "choice", autospec=True) as choice_mock:
      choice_mock.return_value = jnp.zeros((prompt_length), dtype=jnp.int32)
      rng = jax.random.PRNGKey(0)
      init(rng, (prompt_length, embed_size))
      choice_mock.assert_called_once_with(rng, vocab_size, (prompt_length,),
                                          False)

  def test_from_sample_of_embeddings_population_too_small(self):
    vocab_size = 100
    embed_size = 10
    population_size = 5
    prompt_length = 15
    embeddings = np.zeros((vocab_size, embed_size))


# 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/prompts_test.py
import os
import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import seqio
from typing import List
from unittest import mock
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning import prompts
from tensorflow.io import gfile

      choice_mock.assert_called_once_with(rng, population_size,
                                          (prompt_length,), False)

  @parameterized.named_parameters([
      dict(testcase_name="None", population_size=None),
      dict(testcase_name="too_big", population_size=1000)
  ])
  def test_from_sample_of_embeddings_population_size_mismatch(
      self, population_size):
    vocab_size = 100
    embed_size = 10
    prompt_length = 15
    embeddings = np.zeros((vocab_size, embed_size))

    init = prompts.from_sample_of_embeddings(embeddings, population_size)
    with mock.patch.object(
        prompts.jax.random, "choice", autospec=True) as choice_mock:
      choice_mock.return_value = jnp.zeros((prompt_length), dtype=jnp.int32)
      rng = jax.random.PRNGKey(0)
      init(rng, (prompt_length, embed_size))
      choice_mock.assert_called_once_with(rng, vocab_size, (prompt_length,),
                                          False)

  def test_from_sample_of_embeddings_population_too_small(self):
    vocab_size = 100
    embed_size = 10
    population_size = 5
    prompt_length = 15
    embeddings = np.zeros((vocab_size, embed_size))


    init = prompts.from_sample_of_embeddings(embeddings, population_size)