"""simple docstring"""



import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
   import jax
   import jax.numpy as jnp

   from transformers.generation import (
       FlaxForcedBOSTokenLogitsProcessor,
       FlaxForcedEOSTokenLogitsProcessor,
       FlaxLogitsProcessorList,
       FlaxMinLengthLogitsProcessor,
       FlaxTemperatureLogitsWarper,
       FlaxTopKLogitsWarper,
       FlaxTopPLogitsWarper,
   )






@require_flax
class  snake_case__      (       unittest.TestCase       ):


 def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase						):
        __a       =    jnp.ones((batch_size, length)						) / length
        return scores


 def 						a__   ( self						):
        __a       =    None
        __a       =    20

        __a       =    self._get_uniform_logits(batch_size=2    ,		length=lowerCamelCase						)

        # tweak scores to not be uniform anymore
        __a       =    scores.at[1, 5].set((1 / length) + 0.1						)  # peak, 1st batch
        __a       =    scores.at[1, 10].set((1 / length) - 0.4						)  # valley, 1st batch

        # compute softmax
        __a       =    jax.nn.softmax(lowerCamelCase    ,		axis=-1						)

        __a       =    FlaxTemperatureLogitsWarper(temperature=0.5						)
        __a       =    FlaxTemperatureLogitsWarper(temperature=1.3						)

        __a       =    jax.nn.softmax(temp_dist_warper_sharper(lowerCamelCase    ,		scores.copy()    ,		cur_len=lowerCamelCase						)    ,		axis=-1						)
        __a       =    jax.nn.softmax(temp_dist_warper_smoother(lowerCamelCase    ,		scores.copy()    ,		cur_len=lowerCamelCase						)    ,		axis=-1						)

        # uniform distribution stays uniform
        self.assertTrue(jnp.allclose(probs[0, :]    ,		warped_prob_sharp[0, :]    ,		atol=1E-3						)						)
        self.assertTrue(jnp.allclose(probs[0, :]    ,		warped_prob_smooth[0, :]    ,		atol=1E-3						)						)

        # sharp peaks get higher, valleys get lower
        self.assertLess(probs[1, :].max()    ,		warped_prob_sharp[1, :].max()						)
        self.assertGreater(probs[1, :].min()    ,		warped_prob_sharp[1, :].min()						)

        # smooth peaks get lower, valleys get higher
        self.assertGreater(probs[1, :].max()    ,		warped_prob_smooth[1, :].max()						)
        self.assertLess(probs[1, :].min()    ,		warped_prob_smooth[1, :].min()						)


 def 						a__   ( self						):
        __a       =    None
        __a       =    10
        __a       =    2

        # create ramp distribution
        __a       =    np.broadcast_to(np.arange(lowerCamelCase						)[None, :]    ,		(batch_size, vocab_size)						).copy()
        __a       =    ramp_logits[1:, : vocab_size // 2] + vocab_size

        __a       =    FlaxTopKLogitsWarper(3						)

        __a       =    top_k_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)

        # check that correct tokens are filtered
        self.assertListEqual(jnp.isinf(scores[0]						).tolist()    ,		7 * [True] + 3 * [False]						)
        self.assertListEqual(jnp.isinf(scores[1]						).tolist()    ,		2 * [True] + 3 * [False] + 5 * [True]						)

        # check special case
        __a       =    5
        __a       =    FlaxTopKLogitsWarper(top_k=1    ,		filter_value=0.0    ,		min_tokens_to_keep=3						)

        __a       =    np.broadcast_to(np.arange(lowerCamelCase						)[None, :]    ,		(batch_size, length)						).copy()
        __a       =    top_k_warp_safety_check(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)

        # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
        self.assertListEqual((scores == 0.0).sum(axis=-1						).tolist()    ,		[2, 2]						)


 def 						a__   ( self						):
        __a       =    None
        __a       =    10
        __a       =    2

        # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
        __a       =    np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]						)						)

        __a       =    FlaxTopPLogitsWarper(0.8						)
        __a       =    np.exp(top_p_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)						)

        # dist should be filtered to keep min num values so that sum is >= top_p
        # exp (-inf) => 0
        __a       =    np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]						)
        self.assertTrue(np.allclose(lowerCamelCase    ,		lowerCamelCase    ,		atol=1E-3						)						)

        # check edge cases with negative and extreme logits
        __a       =    np.broadcast_to(np.arange(lowerCamelCase						)[None, :]    ,		(batch_size, vocab_size)						).copy() - (
            vocab_size // 2
        )

        # make ramp_logits more extreme
        __a       =    ramp_logits[1] * 100.0

        # make sure at least 2 tokens are kept
        __a       =    FlaxTopPLogitsWarper(0.9    ,		min_tokens_to_keep=2    ,		filter_value=0.0						)
        __a       =    top_p_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)

        # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
        self.assertListEqual((filtered_dist != 0.0).sum(axis=-1						).tolist()    ,		[3, 2]						)


 def 						a__   ( self						):
        __a       =    20
        __a       =    4
        __a       =    0

        __a       =    FlaxMinLengthLogitsProcessor(min_length=10    ,		eos_token_id=lowerCamelCase						)

        # check that min length is applied at length 5
        __a       =    ids_tensor((batch_size, 20)    ,		vocab_size=20						)
        __a       =    5
        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    min_dist_processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()    ,		4 * [-float("inf"						)]						)

        # check that min length is not applied anymore at length 15
        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    15
        __a       =    min_dist_processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        self.assertFalse(jnp.isinf(lowerCamelCase						).any()						)


 def 						a__   ( self						):
        __a       =    20
        __a       =    4
        __a       =    0

        __a       =    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase						)

        # check that all scores are -inf except the bos_token_id score
        __a       =    ids_tensor((batch_size, 1)    ,		vocab_size=20						)
        __a       =    1
        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    logits_processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]						).all()						)
        self.assertListEqual(scores[:, bos_token_id].tolist()    ,		4 * [0]						)  # score for bos_token_id shold be zero

        # check that bos_token_id is not forced if current length is greater than 1
        __a       =    3
        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    logits_processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        self.assertFalse(jnp.isinf(lowerCamelCase						).any()						)


 def 						a__   ( self						):
        __a       =    20
        __a       =    4
        __a       =    0
        __a       =    5

        __a       =    FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase    ,		eos_token_id=lowerCamelCase						)

        # check that all scores are -inf except the eos_token_id when max_length is reached
        __a       =    ids_tensor((batch_size, 4)    ,		vocab_size=20						)
        __a       =    4
        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    logits_processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]						).all()						)
        self.assertListEqual(scores[:, eos_token_id].tolist()    ,		4 * [0]						)  # score for eos_token_id should be zero

        # check that eos_token_id is not forced if max_length is not reached
        __a       =    3
        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    logits_processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        self.assertFalse(jnp.isinf(lowerCamelCase						).any()						)


 def 						a__   ( self						):
        __a       =    4
        __a       =    10
        __a       =    15
        __a       =    2
        __a       =    1
        __a       =    15

        # dummy input_ids and scores
        __a       =    ids_tensor((batch_size, sequence_length)    ,		lowerCamelCase						)
        __a       =    input_ids.copy()

        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    scores.copy()

        # instantiate all dist processors
        __a       =    FlaxTemperatureLogitsWarper(temperature=0.5						)
        __a       =    FlaxTopKLogitsWarper(3						)
        __a       =    FlaxTopPLogitsWarper(0.8						)

        # instantiate all logits processors
        __a       =    FlaxMinLengthLogitsProcessor(min_length=10    ,		eos_token_id=lowerCamelCase						)
        __a       =    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase						)
        __a       =    FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase    ,		eos_token_id=lowerCamelCase						)

        __a       =    10

        # no processor list
        __a       =    temp_dist_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        __a       =    top_k_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        __a       =    top_p_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        __a       =    min_dist_proc(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        __a       =    bos_dist_proc(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
        __a       =    eos_dist_proc(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)

        # with processor list
        __a       =    FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
        __a       =    processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)

        # scores should be equal
        self.assertTrue(jnp.allclose(lowerCamelCase    ,		lowerCamelCase    ,		atol=1E-3						)						)

        # input_ids should never be changed
        self.assertListEqual(input_ids.tolist()    ,		input_ids_comp.tolist()						)






 def 						a__   ( self						):
        __a       =    4
        __a       =    10
        __a       =    15
        __a       =    2
        __a       =    1
        __a       =    15

        # dummy input_ids and scores
        __a       =    ids_tensor((batch_size, sequence_length)    ,		lowerCamelCase						)
        __a       =    input_ids.copy()

        __a       =    self._get_uniform_logits(lowerCamelCase    ,		lowerCamelCase						)
        __a       =    scores.copy()

        # instantiate all dist processors
        __a       =    FlaxTemperatureLogitsWarper(temperature=0.5						)
        __a       =    FlaxTopKLogitsWarper(3						)
        __a       =    FlaxTopPLogitsWarper(0.8						)

        # instantiate all logits processors
        __a       =    FlaxMinLengthLogitsProcessor(min_length=10    ,		eos_token_id=lowerCamelCase						)
        __a       =    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase						)
        __a       =    FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase    ,		eos_token_id=lowerCamelCase						)

        __a       =    10

        # no processor list
        def run_no_processor_list(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
               __a       =    temp_dist_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               __a       =    top_k_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               __a       =    top_p_warp(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               __a       =    min_dist_proc(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               __a       =    bos_dist_proc(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               __a       =    eos_dist_proc(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               return scores

        # with processor list
        def run_processor_list(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
               __a       =    FlaxLogitsProcessorList(
                   [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
               __a       =    processor(lowerCamelCase    ,		lowerCamelCase    ,		cur_len=lowerCamelCase						)
               return scores

        __a       =    jax.jit(lowerCamelCase						)
        __a       =    jax.jit(lowerCamelCase						)

        __a       =    jitted_run_no_processor_list(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						)
        __a       =    jitted_run_processor_list(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						)

        # scores should be equal
        self.assertTrue(jnp.allclose(lowerCamelCase    ,		lowerCamelCase    ,		atol=1E-3						)						)

        # input_ids should never be changed
        self.assertListEqual(input_ids.tolist()    ,		input_ids_comp.tolist()						)

"""simple docstring"""



import unittest

import torch
from torch import nn

from accelerate.test_utils import require_cuda
from accelerate.utils.memory import find_executable_batch_size, release_memory



def 	_lowerCamelCase(	):
       raise RuntimeError("CUDA out of memory."				)






class  snake_case__      (       nn.Module       ):


 def __init__( self						):
        super().__init__()
        __a       =    nn.Linear(3    ,		4						)
        __a       =    nn.BatchNormad(4						)
        __a       =    nn.Linear(4    ,		5						)






 def 						a__   ( self    ,		lowerCamelCase						):
        return self.lineara(self.batchnorm(self.lineara(lowerCamelCase						)						)						)






class  snake_case__      (       unittest.TestCase       ):


 def 						a__   ( self						):
        __a       =    []

        @find_executable_batch_size(starting_batch_size=128						)
        def mock_training_loop_function(lowerCamelCase						):
               nonlocal batch_sizes
               batch_sizes.append(lowerCamelCase						)
               if batch_size != 8:
                      raise_fake_out_of_memory()

        mock_training_loop_function()
        self.assertListEqual(lowerCamelCase    ,		[128, 64, 32, 16, 8]						)


 def 						a__   ( self						):
        __a       =    []

        @find_executable_batch_size(starting_batch_size=128						)
        def mock_training_loop_function(lowerCamelCase    ,		lowerCamelCase						):
               nonlocal batch_sizes
               batch_sizes.append(lowerCamelCase						)
               if batch_size != 8:
                      raise_fake_out_of_memory()
               return batch_size, arga

        __a ,   __a       =    mock_training_loop_function("hello"						)
        self.assertListEqual(lowerCamelCase    ,		[128, 64, 32, 16, 8]						)
        self.assertListEqual([bs, arga]    ,		[8, "hello"]						)


 def 						a__   ( self						):
        @find_executable_batch_size(starting_batch_size=0						)
        def mock_training_loop_function(lowerCamelCase						):
               pass

        with self.assertRaises(lowerCamelCase						) as cm:
               mock_training_loop_function()
               self.assertIn("No executable batch size found, reached zero."    ,		cm.exception.args[0]						)


 def 						a__   ( self						):
        @find_executable_batch_size(starting_batch_size=16						)
        def mock_training_loop_function(lowerCamelCase						):
               if batch_size > 0:
                      raise_fake_out_of_memory()
               pass

        with self.assertRaises(lowerCamelCase						) as cm:
               mock_training_loop_function()
               self.assertIn("No executable batch size found, reached zero."    ,		cm.exception.args[0]						)


 def 						a__   ( self						):
        @find_executable_batch_size(starting_batch_size=128						)
        def mock_training_loop_function(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
               if batch_size != 8:
                      raise raise_fake_out_of_memory()

        with self.assertRaises(lowerCamelCase						) as cm:
               mock_training_loop_function(128    ,		"hello"    ,		"world"						)
               self.assertIn("Batch size was passed into `f`"    ,		cm.exception.args[0]						)
               self.assertIn("`f(arg1='hello', arg2='world')"    ,		cm.exception.args[0]						)


 def 						a__   ( self						):
        @find_executable_batch_size(starting_batch_size=16						)
        def mock_training_loop_function(lowerCamelCase						):
               raise ValueError("Oops, we had an error!"						)

        with self.assertRaises(lowerCamelCase						) as cm:
               mock_training_loop_function()
               self.assertIn("Oops, we had an error!"    ,		cm.exception.args[0]						)






 @require_cuda
 def 						a__   ( self						):
        __a       =    torch.cuda.memory_allocated()
        __a       =    ModelForTest()
        model.cuda()
        self.assertGreater(torch.cuda.memory_allocated()    ,		lowerCamelCase						)
        __a       =    release_memory(lowerCamelCase						)
        self.assertEqual(torch.cuda.memory_allocated()    ,		lowerCamelCase						)
