# 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
        prompt_length,
        set(("prompt->prompt;causal",)),
        multitask=multitask)
    mask = jax.jit(mask_fn)(targets, causal)

    if multitask:
      example_lengths = example_lengths - 1

    for b in range(mask.shape[0]):
      example_length = example_lengths[b] + prompt_length
      for q in range(mask.shape[2]):
        for k in range(mask.shape[3]):
          if q >= example_length or k >= example_length:
            self.assertEqual(mask[b, 1, q, k], 0)
          elif k <= q:
            self.assertEqual(mask[b, 1, q, k], 1)
          else:
            self.assertEqual(mask[b, 1, q, k], 0)

  @parameterized.named_parameters([
      dict(testcase_name="multitask", multitask=True),
      dict(testcase_name="single_task", multitask=False)
  ])
  def test_prompt_decoder_attention_input_visibility(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

        prompt_length,
        set(("prompt->prompt;causal",)),
        multitask=multitask)
    mask = jax.jit(mask_fn)(targets, causal)

    if multitask:
      example_lengths = example_lengths - 1

    for b in range(mask.shape[0]):
      example_length = example_lengths[b] + prompt_length
      for q in range(mask.shape[2]):
        for k in range(mask.shape[3]):
          if q >= example_length or k >= example_length:
            self.assertEqual(mask[b, 1, q, k], 0)
          elif k <= q:
            self.assertEqual(mask[b, 1, q, k], 1)
          else:
            self.assertEqual(mask[b, 1, q, k], 0)

  @parameterized.named_parameters([
      dict(testcase_name="multitask", multitask=True),
      dict(testcase_name="single_task", multitask=False)
  ])
  def test_prompt_decoder_attention_input_visibility(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(