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


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning.extended import multitask_prompts as prompts
import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np

  def test_multi_task_lookup_tasks(self):
    num_tasks = 8
    length = 5
    task_size = 3
    gold_tasks = np.tile(
        np.reshape(np.arange(num_tasks), (num_tasks, 1)), (1, task_size))
    prompt = prompts.MultiTaskPrompt(
        length=length,
        num_tasks=num_tasks,
        shared_init=nn.initializers.zeros,
        tasks_init=lambda *args, **kwargs: gold_tasks,
        get_shared_features=lambda _, x: x,
        get_task_size=lambda _, x, y: task_size,
        combine_shared_and_tasks=lambda _, tasks: tasks,
    )
    params = prompt.init(
        jax.random.PRNGKey(0), jnp.zeros((2, 3), dtype=jnp.int32),
        jnp.zeros((2, 3, 3), dtype=jnp.int32))
    input_ = jnp.arange(num_tasks)
    input_ = jax.random.permutation(jax.random.PRNGKey(0),
                                    input_,
                                    independent=True)
    tasks = jax.jit(prompt.apply)(params, jnp.reshape(input_, (-1, 1)),
                                  jnp.zeros((num_tasks, 1, task_size)))
    for task, idx in zip(tasks, input_):
      np.testing.assert_array_equal(task, idx)

  def test_input_features(self):
    x_embed_features = 512

# 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/multitask_prompts_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning.extended import multitask_prompts as prompts
import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np


  def test_multi_task_lookup_tasks(self):
    num_tasks = 8
    length = 5
    task_size = 3
    gold_tasks = np.tile(
        np.reshape(np.arange(num_tasks), (num_tasks, 1)), (1, task_size))
    prompt = prompts.MultiTaskPrompt(
        length=length,
        num_tasks=num_tasks,
        shared_init=nn.initializers.zeros,
        tasks_init=lambda *args, **kwargs: gold_tasks,
        get_shared_features=lambda _, x: x,
        get_task_size=lambda _, x, y: task_size,
        combine_shared_and_tasks=lambda _, tasks: tasks,
    )
    params = prompt.init(
        jax.random.PRNGKey(0), jnp.zeros((2, 3), dtype=jnp.int32),
        jnp.zeros((2, 3, 3), dtype=jnp.int32))
    input_ = jnp.arange(num_tasks)
    input_ = jax.random.permutation(jax.random.PRNGKey(0),
                                    input_,
                                    independent=True)
    tasks = jax.jit(prompt.apply)(params, jnp.reshape(input_, (-1, 1)),
                                  jnp.zeros((num_tasks, 1, task_size)))
    for task, idx in zip(tasks, input_):
      np.testing.assert_array_equal(task, idx)

  def test_input_features(self):
    x_embed_features = 512

    shared_features = prompts.input_features(None, x_embed_features)