# 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

"""Tests for prompt_tuning.masks."""



def create_input(lengths):
  input_tokens = jnp.ones((len(lengths), jnp.max(lengths)), dtype=jnp.float32)
  input_tokens = input_tokens * (
      jnp.arange(jnp.max(lengths)) < jnp.reshape(lengths, (-1, 1)))
  return input_tokens, lengths


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 CreatePromptDecoderOnlyMaskTest(parameterized.TestCase):

  @parameterized.parameters((jnp.float32,), (jnp.bfloat16,))
  def test_dtype(self, dtype):
    prompt_length = 5
    targets, causal, _ = create_decoder_input(
        jnp.array([2, 3, 4]), jnp.array([5, 6, 8]))

# 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


"""Tests for prompt_tuning.masks."""



def create_input(lengths):
  input_tokens = jnp.ones((len(lengths), jnp.max(lengths)), dtype=jnp.float32)
  input_tokens = input_tokens * (
      jnp.arange(jnp.max(lengths)) < jnp.reshape(lengths, (-1, 1)))
  return input_tokens, lengths


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 CreatePromptDecoderOnlyMaskTest(parameterized.TestCase):

  @parameterized.parameters((jnp.float32,), (jnp.bfloat16,))
  def test_dtype(self, dtype):
    prompt_length = 5
    targets, causal, _ = create_decoder_input(
        jnp.array([2, 3, 4]), jnp.array([5, 6, 8]))

    mask_fn = masks.create_prompt_decoder_only_mask(prompt_length)