



lowerCAmelCase__					  =  '''0.21.0'''

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
				from .utils import rich









import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
				import optax
				from flax.training.common_utils import onehot

				from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
				from transformers.models.ta.modeling_flax_ta import shift_tokens_right

@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 				snake_case__(unittest.TestCase      ):

					"""simple docstring"""






					@slow
					def 			snake_case					(							self :     Optional[int]					):
									lowercase__					:   Dict													=					FlaxMTaForConditionalGeneration.from_pretrained("google/mt5-small"					)
									lowercase__					:   Dict													=					AutoTokenizer.from_pretrained("google/mt5-small"					)

									lowercase__					:   Optional[Any]													=					tokenizer("Hello there"				,       return_tensors="np"					).input_ids
									lowercase__					:   Optional[Any]													=					tokenizer("Hi I am"				,       return_tensors="np"					).input_ids

									lowercase__					:   int													=					shift_tokens_right(SCREAMING_SNAKE_CASE				,       model.config.pad_token_id				,       model.config.decoder_start_token_id					)

									lowercase__					:   Tuple													=					model(SCREAMING_SNAKE_CASE				,       decoder_input_ids=SCREAMING_SNAKE_CASE					).logits
									lowercase__					:   Dict													=					optax.softmax_cross_entropy(SCREAMING_SNAKE_CASE				,       onehot(SCREAMING_SNAKE_CASE				,       logits.shape[-1]					)					).mean()

									lowercase__					:   Union[str, Any]													=					-(labels.shape[-1] * loss.item())

									lowercase__					:   Union[str, Any]													=					-84.9_127
									self.assertTrue(abs(mtf_score - EXPECTED_SCORE					) < 1E-4					)




