





import unittest

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


if is_flax_available():
     import jax.numpy as jnp

     from transformers import FlaxXLMRobertaModel

@require_sentencepiece
@require_tokenizers
@require_flax
class      A       (	unittest.TestCase						):



      @slow
      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    str:
         '''simple docstring'''

         UpperCAmelCase_											=FlaxXLMRobertaModel.from_pretrained("xlm-roberta-base"       )
         UpperCAmelCase_											=AutoTokenizer.from_pretrained("xlm-roberta-base"       )
         UpperCAmelCase_											="The dog is cute and lives in the garden house"
         UpperCAmelCase_											=jnp.array([tokenizer.encode(_lowerCAmelCase       )]       )

         UpperCAmelCase_											=(1, 12, 768)  # batch_size, sequence_length, embedding_vector_dim
         UpperCAmelCase_											=jnp.array(
             [[-0.01_01, 0.12_18, -0.08_03, 0.08_01, 0.13_27, 0.07_76, -0.12_15, 0.23_83, 0.33_38, 0.31_06, 0.03_00, 0.02_52]]       )

         UpperCAmelCase_											=model(_lowerCAmelCase       )["last_hidden_state"]
         self.assertEqual(output.shape		,  _lowerCAmelCase       )
         # compare the actual values for a slice of last dim
         self.assertTrue(jnp.allclose(output[:, :, -1]		,  _lowerCAmelCase		,  atol=1e-3       )       )











__lowercase  :   Dict						      ="""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




