# Path: t5x/decoding.py
# NEG_INF = np.array(-1.0e7)
# def temperature_sample(
#     inputs: jnp.ndarray,
#     cache: Mapping[str, jnp.ndarray],
#     tokens_to_logits: Callable[[jnp.ndarray, Mapping[str, jnp.ndarray]],
#                                Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]],
#     eos_id: int,
#     decode_rng: Optional[jnp.ndarray] = None,
#     num_decodes: int = 1,
#     temperature: float = 1.0,
#     topk: int = 1,
#     topp: float = 0.0,
#     cache_offset: int = 0,
#     initial_index: Optional[jnp.ndarray] = None,
#     max_decode_steps: Optional[int] = None,
# ) -> Tuple[jnp.ndarray, jnp.ndarray]:
# def _temperature_sample_single_trial(
#     inputs: jnp.ndarray,
#     cache: Mapping[str, jnp.ndarray],
#     tokens_to_logits: Callable[[jnp.ndarray, Mapping[str, jnp.ndarray]],
#                                Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]],
#     eos_id: int,
#     prng_key: jnp.ndarray,
#     temperature: float = 1.0,
#     topk: int = 20,
#     topp: float = 0.0,
#     initial_index: Optional[jnp.ndarray] = None,
#     max_decode_steps: Optional[int] = None) -> jnp.ndarray:
#   def sampling_loop_cond_fn(state: SamplingLoopState) -> bool:
#   def sampling_loop_body_fn(state: SamplingLoopState) -> SamplingLoopState:
# def brevity_penalty(alpha: float, length: int) -> jnp.ndarray:
# def cache_map(fn, cache, apply_to_index: bool = False):
# def add_beam_dim(x: jnp.ndarray,
#                  beam_size: int,
#                  offset: int = 0) -> jnp.ndarray:
# def flatten_beam_dim(x: jnp.ndarray, offset: int = 0) -> jnp.ndarray:
# def unflatten_beam_dim(x: jnp.ndarray,
#                        batch_size: int,
#                        beam_size: int,
#                        offset: int = 0) -> jnp.ndarray:
# def flat_batch_beam_expand(x: jnp.ndarray,
#                            beam_size: int,
#                            offset: int = 0) -> jnp.ndarray:
# def cache_gather_beams(nested: PyTreeDef,
#                        beam_indices: jnp.ndarray,
#                        batch_size: int,
#                        old_beam_size: int,
#                        new_beam_size: int,
#                        one_hot: bool = True,
#                        offset: int = 0) -> jnp.ndarray:
#       def gather_fn(x):
#       def gather_fn(x):
#       def gather_fn(x):
#       def gather_fn(x):
# def gather_beams(nested: PyTreeDef,
#                  beam_indices: jnp.ndarray,
#                  batch_size: int,
#                  old_beam_size: int,
#                  new_beam_size: int,
#                  one_hot: bool = True) -> jnp.ndarray:
#     def gather_fn(x):
#     def gather_fn(x):
# def top_k_two_stage(x, k):
# def gather_topk_beams(nested: PyTreeDef, score_or_log_prob: jnp.ndarray,
#                       batch_size: int, new_beam_size: int) -> jnp.ndarray:
# def beam_init(batch_size: int,
#               beam_size: int,
#               max_decode_len: int,
#               cache: Mapping[str, jnp.ndarray],
#               offset: int = 0) -> BeamState:
# def beam_search(inputs: jnp.ndarray,
#                 cache: Mapping[str, jnp.ndarray],
#                 tokens_to_logits: Callable[
#                     [jnp.ndarray, Mapping[str, jnp.ndarray]],
#                     Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]],
#                 eos_id: int,
#                 num_decodes: int = 4,
#                 alpha: float = 0.6,
#                 max_decode_len: Optional[int] = None,
#                 decode_rng: Optional[jnp.ndarray] = None,
#                 cache_offset: int = 0) -> Tuple[jnp.ndarray, jnp.ndarray]:
#   def beam_search_loop_cond_fn(state: BeamState) -> bool:
#   def beam_search_loop_body_fn(state: BeamState) -> BeamState:
# class BeamState:


import functools
import jax
import jax.numpy as jnp
import numpy as np
from typing import Mapping, Tuple
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from t5x import decoding
    inputs = np.zeros((1, 1), dtype=np.int32)

    token_to_logits = mock.Mock()

    # logits correspond to (0.3, 0, 0.1, 0.6)
    token_to_logits.return_value = (np.array([[-1.2, NEG_INF, -2.3, -0.51]],
                                             dtype=np.float32), {})

    with jax.disable_jit():
      # this lets us see logits after topp and topk are applied
      with mock.patch.object(jax.random, 'categorical') as mocked:
        mocked.return_value = jnp.array([0], dtype=jnp.int32)
        decodes, _ = decoding.temperature_sample(
            inputs, {},
            token_to_logits,
            EOS_ID,
            rng0,
            temperature=1.4,
            topp=0.7,
            topk=0)

    self.assertLen(token_to_logits.call_args_list, 1)
    np.testing.assert_array_equal(decodes, jnp.asarray([[[0]]]))

    np.testing.assert_array_almost_equal(
        mocked.call_args_list[0][0][1],
        jnp.asarray([[-0.85714293, NEG_INF, NEG_INF, -0.36428571]]))

  def test_add_beam_dim(self):
    x = np.array([[0, 5, 1, 0], [0, 8, 6, 9]], dtype=np.int32)

# Path: t5x/decoding.py
# NEG_INF = np.array(-1.0e7)
# def temperature_sample(
#     inputs: jnp.ndarray,
#     cache: Mapping[str, jnp.ndarray],
#     tokens_to_logits: Callable[[jnp.ndarray, Mapping[str, jnp.ndarray]],
#                                Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]],
#     eos_id: int,
#     decode_rng: Optional[jnp.ndarray] = None,
#     num_decodes: int = 1,
#     temperature: float = 1.0,
#     topk: int = 1,
#     topp: float = 0.0,
#     cache_offset: int = 0,
#     initial_index: Optional[jnp.ndarray] = None,
#     max_decode_steps: Optional[int] = None,
# ) -> Tuple[jnp.ndarray, jnp.ndarray]:
# def _temperature_sample_single_trial(
#     inputs: jnp.ndarray,
#     cache: Mapping[str, jnp.ndarray],
#     tokens_to_logits: Callable[[jnp.ndarray, Mapping[str, jnp.ndarray]],
#                                Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]],
#     eos_id: int,
#     prng_key: jnp.ndarray,
#     temperature: float = 1.0,
#     topk: int = 20,
#     topp: float = 0.0,
#     initial_index: Optional[jnp.ndarray] = None,
#     max_decode_steps: Optional[int] = None) -> jnp.ndarray:
#   def sampling_loop_cond_fn(state: SamplingLoopState) -> bool:
#   def sampling_loop_body_fn(state: SamplingLoopState) -> SamplingLoopState:
# def brevity_penalty(alpha: float, length: int) -> jnp.ndarray:
# def cache_map(fn, cache, apply_to_index: bool = False):
# def add_beam_dim(x: jnp.ndarray,
#                  beam_size: int,
#                  offset: int = 0) -> jnp.ndarray:
# def flatten_beam_dim(x: jnp.ndarray, offset: int = 0) -> jnp.ndarray:
# def unflatten_beam_dim(x: jnp.ndarray,
#                        batch_size: int,
#                        beam_size: int,
#                        offset: int = 0) -> jnp.ndarray:
# def flat_batch_beam_expand(x: jnp.ndarray,
#                            beam_size: int,
#                            offset: int = 0) -> jnp.ndarray:
# def cache_gather_beams(nested: PyTreeDef,
#                        beam_indices: jnp.ndarray,
#                        batch_size: int,
#                        old_beam_size: int,
#                        new_beam_size: int,
#                        one_hot: bool = True,
#                        offset: int = 0) -> jnp.ndarray:
#       def gather_fn(x):
#       def gather_fn(x):
#       def gather_fn(x):
#       def gather_fn(x):
# def gather_beams(nested: PyTreeDef,
#                  beam_indices: jnp.ndarray,
#                  batch_size: int,
#                  old_beam_size: int,
#                  new_beam_size: int,
#                  one_hot: bool = True) -> jnp.ndarray:
#     def gather_fn(x):
#     def gather_fn(x):
# def top_k_two_stage(x, k):
# def gather_topk_beams(nested: PyTreeDef, score_or_log_prob: jnp.ndarray,
#                       batch_size: int, new_beam_size: int) -> jnp.ndarray:
# def beam_init(batch_size: int,
#               beam_size: int,
#               max_decode_len: int,
#               cache: Mapping[str, jnp.ndarray],
#               offset: int = 0) -> BeamState:
# def beam_search(inputs: jnp.ndarray,
#                 cache: Mapping[str, jnp.ndarray],
#                 tokens_to_logits: Callable[
#                     [jnp.ndarray, Mapping[str, jnp.ndarray]],
#                     Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]],
#                 eos_id: int,
#                 num_decodes: int = 4,
#                 alpha: float = 0.6,
#                 max_decode_len: Optional[int] = None,
#                 decode_rng: Optional[jnp.ndarray] = None,
#                 cache_offset: int = 0) -> Tuple[jnp.ndarray, jnp.ndarray]:
#   def beam_search_loop_cond_fn(state: BeamState) -> bool:
#   def beam_search_loop_body_fn(state: BeamState) -> BeamState:
# class BeamState:


# Path: t5x/decoding_test.py
import functools
import jax
import jax.numpy as jnp
import numpy as np
from typing import Mapping, Tuple
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from t5x import decoding

    inputs = np.zeros((1, 1), dtype=np.int32)

    token_to_logits = mock.Mock()

    # logits correspond to (0.3, 0, 0.1, 0.6)
    token_to_logits.return_value = (np.array([[-1.2, NEG_INF, -2.3, -0.51]],
                                             dtype=np.float32), {})

    with jax.disable_jit():
      # this lets us see logits after topp and topk are applied
      with mock.patch.object(jax.random, 'categorical') as mocked:
        mocked.return_value = jnp.array([0], dtype=jnp.int32)
        decodes, _ = decoding.temperature_sample(
            inputs, {},
            token_to_logits,
            EOS_ID,
            rng0,
            temperature=1.4,
            topp=0.7,
            topk=0)

    self.assertLen(token_to_logits.call_args_list, 1)
    np.testing.assert_array_equal(decodes, jnp.asarray([[[0]]]))

    np.testing.assert_array_almost_equal(
        mocked.call_args_list[0][0][1],
        jnp.asarray([[-0.85714293, NEG_INF, NEG_INF, -0.36428571]]))

  def test_add_beam_dim(self):
    x = np.array([[0, 5, 1, 0], [0, 8, 6, 9]], dtype=np.int32)

    y = decoding.add_beam_dim(x, beam_size=3)