# Path: prompt_tuning/extended/multitask_prompts.py
# def identity(x: Any) -> Any:
# def sentinel_to_task(x: Array, vocab_size: int) -> Array:
# def add_shared_and_tasks(shared: Array, tasks: Array) -> Array:
# def concat_shared_and_tasks(shared: Array, tasks: Array) -> Array:
# def concat_features_shared_and_tasks(shared: Array, tasks: Array) -> Array:
# def project_shared_with_tasks(shared: Array, tasks: Array) -> Array:
# def input_features(self, x_embed_features: int) -> int:  # pylint: disable=unused-argument
# def shared_features_attr_or_input_features(self, x_embed_features: int) -> int:
# def match_shared(self, shared_length: int, shared_features: int) -> int:  # pylint: disable=unused-argument
# def task_length_attr_or_match_shared(self, shared_length: int,
#                                      shared_features: int) -> int:
# def task_features_or_match_shared(self, shared_length: int,
#                                   shared_features: int) -> int:
# def projection_size(self, shared_length: int, shared_features: int) -> int:  # pylint: disable=unused-argument
#   def setup(self):
#   def __call__(self, x, x_embed):
# class MultiTaskPrompt(nn.Module):
# class AddMultiTaskPrompt(MultiTaskPrompt):
# class ConcatMultiTaskPrompt(MultiTaskPrompt):
# class ConcatFeaturesMultiTaskPrompt(MultiTaskPrompt):
# class ProjectMultiTaskPrompt(MultiTaskPrompt):
# 
# Path: prompt_tuning/extended/train/multitask_prompts.py
# class MultiTaskPrompt(nn.Module):
#   def __call__(self, x, x_embed):
# 
# 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 prompt_tuning.extended import multitask_prompts
from prompt_tuning.extended.train import multitask_prompts as train_multitask_prompts
from prompt_tuning.train import prompts as train_prompts
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 multitask training prompts."""



class PromptsTest(absltest.TestCase):

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

# Path: prompt_tuning/extended/multitask_prompts.py
# def identity(x: Any) -> Any:
# def sentinel_to_task(x: Array, vocab_size: int) -> Array:
# def add_shared_and_tasks(shared: Array, tasks: Array) -> Array:
# def concat_shared_and_tasks(shared: Array, tasks: Array) -> Array:
# def concat_features_shared_and_tasks(shared: Array, tasks: Array) -> Array:
# def project_shared_with_tasks(shared: Array, tasks: Array) -> Array:
# def input_features(self, x_embed_features: int) -> int:  # pylint: disable=unused-argument
# def shared_features_attr_or_input_features(self, x_embed_features: int) -> int:
# def match_shared(self, shared_length: int, shared_features: int) -> int:  # pylint: disable=unused-argument
# def task_length_attr_or_match_shared(self, shared_length: int,
#                                      shared_features: int) -> int:
# def task_features_or_match_shared(self, shared_length: int,
#                                   shared_features: int) -> int:
# def projection_size(self, shared_length: int, shared_features: int) -> int:  # pylint: disable=unused-argument
#   def setup(self):
#   def __call__(self, x, x_embed):
# class MultiTaskPrompt(nn.Module):
# class AddMultiTaskPrompt(MultiTaskPrompt):
# class ConcatMultiTaskPrompt(MultiTaskPrompt):
# class ConcatFeaturesMultiTaskPrompt(MultiTaskPrompt):
# class ProjectMultiTaskPrompt(MultiTaskPrompt):
# 
# Path: prompt_tuning/extended/train/multitask_prompts.py
# class MultiTaskPrompt(nn.Module):
#   def __call__(self, x, x_embed):
# 
# 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/extended/train/multitask_prompts_test.py
from unittest import mock
from absl.testing import absltest
from prompt_tuning.extended import multitask_prompts
from prompt_tuning.extended.train import multitask_prompts as train_multitask_prompts
from prompt_tuning.train import prompts as train_prompts
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 multitask training prompts."""



class PromptsTest(absltest.TestCase):

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

        train_prompts.prefix_prompt, spec_set=True)