# Path: prompt_tuning/extended/masks.py
# INPUT_INPUT = "input->input"
# INPUT_PROMPT = "input->prompt"
# PROMPT_PROMPT = "prompt->prompt"
# PROMPT_PROMPT_CAUSAL = "prompt->prompt;causal"
# PROMPT_INPUT = "prompt->input"
# def summarize_attention(attentions: Set[str], decoder: bool = False) -> str:
# def prompt_encoder_attention_mask(
#     prompt_length: int,
#     allowable_attention: Optional[Set[str]] = None,
#     multitask: bool = False) -> Callable[[Array, DType], Array]:
#   def _prompt_encoder_attention_mask(encoder_input_tokens: Array,
#                                      dtype: DType = jnp.float32) -> Array:
# def prompt_decoder_attention_mask(
#     prompt_length: int,
#     allowable_attention: Optional[Set[str]] = None,
#     multitask: bool = False) -> Callable[[Array, Array, DType], Array]:
#   def _prompt_decoder_attention_mask(
#       decoder_target_tokens: Array,
#       decoder_causal_attention: Optional[Array],
#       dtype: DType = jnp.float32) -> Array:


import functools
import jax
import jax.numpy as jnp
from absl.testing import absltest
from absl.testing import parameterized
from prompt_tuning.extended import masks
def create_decoder_input(prefix_lengths, target_lengths):
  targets, example_lengths = create_input(prefix_lengths + target_lengths)
  causal, _ = create_input(prefix_lengths)
  causal = jnp.concatenate([
      causal,
      jnp.zeros((causal.shape[0], jnp.max(example_lengths) - causal.shape[1]))
  ],
                           axis=1)
  return targets, causal, example_lengths


class PromptDecoderAttentionTest(parameterized.TestCase):
  """This tests the creation of prompt masks for decoder.

  Note:
    This function is used when the prompt is applied to the decoder, it is not
    needed for decoder that acts on a prompted encoder.
  """

  @parameterized.named_parameters([
      dict(testcase_name="multitask", multitask=True),
      dict(testcase_name="single_task", multitask=False)
  ])
  def test_prompt_decoder_attention_causal_targets(self, multitask):
    prefix_lengths = jnp.array([4, 6, 2, 3, 5])
    target_lengths = jnp.array([2, 1, 3, 1, 1])
    prompt_length = 4
    targets, causal, example_lengths = create_decoder_input(
        prefix_lengths, target_lengths)


# Path: prompt_tuning/extended/masks.py
# INPUT_INPUT = "input->input"
# INPUT_PROMPT = "input->prompt"
# PROMPT_PROMPT = "prompt->prompt"
# PROMPT_PROMPT_CAUSAL = "prompt->prompt;causal"
# PROMPT_INPUT = "prompt->input"
# def summarize_attention(attentions: Set[str], decoder: bool = False) -> str:
# def prompt_encoder_attention_mask(
#     prompt_length: int,
#     allowable_attention: Optional[Set[str]] = None,
#     multitask: bool = False) -> Callable[[Array, DType], Array]:
#   def _prompt_encoder_attention_mask(encoder_input_tokens: Array,
#                                      dtype: DType = jnp.float32) -> Array:
# def prompt_decoder_attention_mask(
#     prompt_length: int,
#     allowable_attention: Optional[Set[str]] = None,
#     multitask: bool = False) -> Callable[[Array, Array, DType], Array]:
#   def _prompt_decoder_attention_mask(
#       decoder_target_tokens: Array,
#       decoder_causal_attention: Optional[Array],
#       dtype: DType = jnp.float32) -> Array:


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

def create_decoder_input(prefix_lengths, target_lengths):
  targets, example_lengths = create_input(prefix_lengths + target_lengths)
  causal, _ = create_input(prefix_lengths)
  causal = jnp.concatenate([
      causal,
      jnp.zeros((causal.shape[0], jnp.max(example_lengths) - causal.shape[1]))
  ],
                           axis=1)
  return targets, causal, example_lengths


class PromptDecoderAttentionTest(parameterized.TestCase):
  """This tests the creation of prompt masks for decoder.

  Note:
    This function is used when the prompt is applied to the decoder, it is not
    needed for decoder that acts on a prompted encoder.
  """

  @parameterized.named_parameters([
      dict(testcase_name="multitask", multitask=True),
      dict(testcase_name="single_task", multitask=False)
  ])
  def test_prompt_decoder_attention_causal_targets(self, multitask):
    prefix_lengths = jnp.array([4, 6, 2, 3, 5])
    target_lengths = jnp.array([2, 1, 3, 1, 1])
    prompt_length = 4
    targets, causal, example_lengths = create_decoder_input(
        prefix_lengths, target_lengths)


    mask_fn = masks.prompt_decoder_attention_mask(