# 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
# Copyright 2022 Google.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for training prompts."""



class PromptsTest(parameterized.TestCase):

  def test_prefix_prompt(self):
    embed_size = 20
    prompt_length = 5
    batch_size = 2
    seq_len = 14
    prompt = jnp.zeros((batch_size, prompt_length, embed_size))
    embed = jnp.ones((batch_size, seq_len, 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/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

# Copyright 2022 Google.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for training prompts."""



class PromptsTest(parameterized.TestCase):

  def test_prefix_prompt(self):
    embed_size = 20
    prompt_length = 5
    batch_size = 2
    seq_len = 14
    prompt = jnp.zeros((batch_size, prompt_length, embed_size))
    embed = jnp.ones((batch_size, seq_len, embed_size))

    with_prompt = jax.jit(train_prompts.prefix_prompt)(prompt, embed, None)