# Path: prompt_tuning/masks.py
# def add_fake_prompt(prompt_length: int,
#                     multitask: bool = False) -> Callable[[Array], Array]:
#   def _add_fake_prompt(encoder_input_tokens: Array):
# def create_prompt_encoder_mask(
#     prompt_length: int) -> Callable[[Array, DType], Array]:
#   def _prompt_encoder_attention_mask(encoder_input_tokens: Array,
#                                      dtype: DType = jnp.float32) -> Array:
# def create_prompt_decoder_only_mask(
#     prompt_length: int) -> Callable[[Array, Array, DType], Array]:
#   def _prompt_decoder_only_mask(decoder_target_tokens: Array,
#                                 decoder_causal_attention: Optional[Array],
#                                 dtype: DType = jnp.float32) -> Array:


import functools
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning import masks
    mask_fn = masks.create_prompt_encoder_mask(prompt_length)
    mask = jax.jit(mask_fn)(inputs)

    batch, heads, query, key = mask.shape
    lengths = prompt_length + example_lengths
    for b in range(batch):
      for h in range(heads):
        for q in range(query):
          for k in range(key):
            if q < lengths[b] and k < lengths[b]:
              self.assertEqual(mask[b, h, q, k].item(), 1)
            else:
              self.assertEqual(mask[b, h, q, k].item(), 0)


class AddFakePromptTest(parameterized.TestCase):
  """Tests that we update prompted encoder inputs for a non-prompted decoder.

  This function can also be used for a prompted encoder that always allows
  fully visible attention in the prompt.
  """

  @parameterized.named_parameters([
      dict(testcase_name="multitask", multitask=True),
      dict(testcase_name="single_task", multitask=False)
  ])
  def test_add_fake_prompt(self, multitask):
    encoder_input_tokens, _ = create_input(jnp.array([5, 6, 7, 6, 1]))
    prompt_length = 4


# Path: prompt_tuning/masks.py
# def add_fake_prompt(prompt_length: int,
#                     multitask: bool = False) -> Callable[[Array], Array]:
#   def _add_fake_prompt(encoder_input_tokens: Array):
# def create_prompt_encoder_mask(
#     prompt_length: int) -> Callable[[Array, DType], Array]:
#   def _prompt_encoder_attention_mask(encoder_input_tokens: Array,
#                                      dtype: DType = jnp.float32) -> Array:
# def create_prompt_decoder_only_mask(
#     prompt_length: int) -> Callable[[Array, Array, DType], Array]:
#   def _prompt_decoder_only_mask(decoder_target_tokens: Array,
#                                 decoder_causal_attention: Optional[Array],
#                                 dtype: DType = jnp.float32) -> Array:


# Path: prompt_tuning/masks_test.py
import functools
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning import masks

    mask_fn = masks.create_prompt_encoder_mask(prompt_length)
    mask = jax.jit(mask_fn)(inputs)

    batch, heads, query, key = mask.shape
    lengths = prompt_length + example_lengths
    for b in range(batch):
      for h in range(heads):
        for q in range(query):
          for k in range(key):
            if q < lengths[b] and k < lengths[b]:
              self.assertEqual(mask[b, h, q, k].item(), 1)
            else:
              self.assertEqual(mask[b, h, q, k].item(), 0)


class AddFakePromptTest(parameterized.TestCase):
  """Tests that we update prompted encoder inputs for a non-prompted decoder.

  This function can also be used for a prompted encoder that always allows
  fully visible attention in the prompt.
  """

  @parameterized.named_parameters([
      dict(testcase_name="multitask", multitask=True),
      dict(testcase_name="single_task", multitask=False)
  ])
  def test_add_fake_prompt(self, multitask):
    encoder_input_tokens, _ = create_input(jnp.array([5, 6, 7, 6, 1]))
    prompt_length = 4


    pretend_fn = masks.add_fake_prompt(prompt_length, multitask)