

"""simple docstring"""






import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
 import os

 import jax.numpy as jnp
 from jax import jit

 from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
 from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

 _lowerCAmelCase						=		"""0.12"""  # assumed parallelism: 8


if is_torch_available():
 import torch

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase=None   ):

    '''simple docstring'''



    if rng is None:
        _lowerCAmelCase       :   List[Any]  =    random.Random()

    _lowerCAmelCase       :   Dict  =    1
    for dim in shape:
        total_dims *= dim

    _lowerCAmelCase       :   Optional[Any]  =    []
    for _ in range(_lowerCamelCase   ):
        values.append(rng.randint(0 ,   vocab_size - 1   )   )

    _lowerCAmelCase       :   int  =    np.array(_lowerCamelCase ,   dtype=jnp.intaa   ).reshape(_lowerCamelCase   )

    return output

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase=None   ):

    '''simple docstring'''



    _lowerCAmelCase       :   Tuple  =    ids_tensor(_lowerCamelCase ,   vocab_size=2 ,   rng=_lowerCamelCase   )
    # make sure that at least one token is attended to for each batch
    _lowerCAmelCase       :   str  =    1
    return attn_mask







@require_flax
class 	__UpperCamelCase    :
  _UpperCAmelCase          =		None
  _UpperCAmelCase          =		()

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase       :   str  =    self.model_tester.prepare_config_and_inputs_for_common()

      # cut to half length & take max batch_size 3
      _lowerCAmelCase       :   Union[str, Any]  =    2
      _lowerCAmelCase       :   Union[str, Any]  =    inputs['input_ids'].shape[-1] // 2
      _lowerCAmelCase       :   Tuple  =    inputs['input_ids'][:max_batch_size, :sequence_length]

      _lowerCAmelCase       :   int  =    jnp.ones_like(_A      )
      _lowerCAmelCase       :   List[Any]  =    attention_mask[:max_batch_size, :sequence_length]

      # generate max 5 tokens
      _lowerCAmelCase       :   List[Any]  =    input_ids.shape[-1] + 5
      if config.eos_token_id is not None and config.pad_token_id is None:
          # hack to allow generate for models such as GPT2 as is done in `generate()`
          _lowerCAmelCase       :   Union[str, Any]  =    config.eos_token_id
      return config, input_ids, attention_mask, max_length

  @is_pt_flax_cross_test
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Tuple  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   Dict  =    False
      _lowerCAmelCase       :   Dict  =    max_length
      _lowerCAmelCase       :   Tuple  =    0

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   Dict  =    model_class(_A      )

          _lowerCAmelCase       :   Any  =    model_class.__name__[4:]  # Skip the "Flax" at the beginning
          _lowerCAmelCase       :   List[str]  =    getattr(_A    ,_A      )
          _lowerCAmelCase       :   Union[str, Any]  =    pt_model_class(_A      ).eval()
          _lowerCAmelCase       :   Tuple  =    load_flax_weights_in_pytorch_model(_A    ,flax_model.params      )

          _lowerCAmelCase       :   Union[str, Any]  =    flax_model.generate(_A      ).sequences
          _lowerCAmelCase       :   str  =    pt_model.generate(torch.tensor(_A    ,dtype=torch.long      )      )

          if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
              _lowerCAmelCase       :   int  =    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

          self.assertListEqual(pt_generation_outputs.numpy().tolist()    ,flax_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Optional[int]  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   str  =    False
      _lowerCAmelCase       :   Optional[Any]  =    max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   Union[str, Any]  =    model_class(_A      )

          _lowerCAmelCase       :   Union[str, Any]  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   Any  =    jit(model.generate      )
          _lowerCAmelCase       :   Dict  =    jit_generate(_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Any  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   Union[str, Any]  =    True
      _lowerCAmelCase       :   Union[str, Any]  =    max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   Tuple  =    model_class(_A      )

          _lowerCAmelCase       :   Any  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   Optional[Any]  =    jit(model.generate      )
          _lowerCAmelCase       :   Tuple  =    jit_generate(_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Tuple  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   str  =    False
      _lowerCAmelCase       :   List[str]  =    max_length
      _lowerCAmelCase       :   int  =    2

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   List[Any]  =    model_class(_A      )

          _lowerCAmelCase       :   Tuple  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   List[str]  =    jit(model.generate      )
          _lowerCAmelCase       :   Optional[Any]  =    jit_generate(_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Union[str, Any]  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   int  =    False
      _lowerCAmelCase       :   Dict  =    max_length
      _lowerCAmelCase       :   Optional[Any]  =    2
      _lowerCAmelCase       :   List[str]  =    2

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   Optional[Any]  =    model_class(_A      )

          _lowerCAmelCase       :   str  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[0]    ,input_ids.shape[0] * config.num_return_sequences      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Any  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   Optional[int]  =    True
      _lowerCAmelCase       :   int  =    max_length
      _lowerCAmelCase       :   str  =    0.8
      _lowerCAmelCase       :   List[Any]  =    10
      _lowerCAmelCase       :   Any  =    0.3
      _lowerCAmelCase       :   int  =    1
      _lowerCAmelCase       :   int  =    8
      _lowerCAmelCase       :   Tuple  =    9

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   List[Any]  =    model_class(_A      )

          _lowerCAmelCase       :   Any  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   Any  =    jit(model.generate      )
          _lowerCAmelCase       :   Optional[int]  =    jit_generate(_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Optional[Any]  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   Dict  =    max_length
      _lowerCAmelCase       :   int  =    1
      _lowerCAmelCase       :   str  =    8
      _lowerCAmelCase       :   Optional[int]  =    9

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   int  =    model_class(_A      )

          _lowerCAmelCase       :   Union[str, Any]  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   str  =    jit(model.generate      )
          _lowerCAmelCase       :   Union[str, Any]  =    jit_generate(_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   int  =    self._get_input_ids_and_config()
      _lowerCAmelCase       :   Optional[Any]  =    max_length
      _lowerCAmelCase       :   Dict  =    2
      _lowerCAmelCase       :   Optional[Any]  =    1
      _lowerCAmelCase       :   List[Any]  =    8
      _lowerCAmelCase       :   List[str]  =    9

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   List[str]  =    model_class(_A      )

          _lowerCAmelCase       :   List[Any]  =    model.generate(_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   str  =    jit(model.generate      )
          _lowerCAmelCase       :   str  =    jit_generate(_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Union[str, Any]  =    self._get_input_ids_and_config()

      # pad attention mask on the left
      _lowerCAmelCase       :   Union[str, Any]  =    attention_mask.at[(0, 0)].set(0      )

      _lowerCAmelCase       :   Optional[Any]  =    False
      _lowerCAmelCase       :   Optional[int]  =    max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   List[str]  =    model_class(_A      )

          _lowerCAmelCase       :   Optional[int]  =    model.generate(_A    ,attention_mask=_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   Optional[int]  =    jit(model.generate      )
          _lowerCAmelCase       :   Optional[Any]  =    jit_generate(_A    ,attention_mask=_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Optional[Any]  =    self._get_input_ids_and_config()

      # pad attention mask on the left
      _lowerCAmelCase       :   Dict  =    attention_mask.at[(0, 0)].set(0      )

      _lowerCAmelCase       :   List[Any]  =    True
      _lowerCAmelCase       :   Any  =    max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   str  =    model_class(_A      )

          _lowerCAmelCase       :   Any  =    model.generate(_A    ,attention_mask=_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   Any  =    jit(model.generate      )
          _lowerCAmelCase       :   Dict  =    jit_generate(_A    ,attention_mask=_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Optional[int]  =    self._get_input_ids_and_config()

      # pad attention mask on the left
      _lowerCAmelCase       :   Optional[int]  =    attention_mask.at[(0, 0)].set(0      )

      _lowerCAmelCase       :   Dict  =    2
      _lowerCAmelCase       :   Any  =    max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase       :   List[Any]  =    model_class(_A      )

          _lowerCAmelCase       :   Union[str, Any]  =    model.generate(_A    ,attention_mask=_A      ).sequences
          self.assertEqual(generation_outputs.shape[-1]    ,_A      )

          _lowerCAmelCase       :   Any  =    jit(model.generate      )
          _lowerCAmelCase       :   Any  =    jit_generate(_A    ,attention_mask=_A      ).sequences

          self.assertListEqual(generation_outputs.tolist()    ,jit_generation_outputs.tolist()      )







@require_flax
class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    AutoTokenizer.from_pretrained('hf-internal-testing/tiny-bert'      )
      _lowerCAmelCase       :   List[Any]  =    FlaxAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-bert-flax-only'      )

      _lowerCAmelCase       :   Any  =    'Hello world'
      _lowerCAmelCase       :   Tuple  =    tokenizer(_A    ,return_tensors='np'      ).input_ids

      # typos are quickly detected (the correct argument is `do_sample`)
      with self.assertRaisesRegex(_A    ,'do_samples'      ):
          model.generate(_A    ,do_samples=_A      )

      # arbitrary arguments that will not be used anywhere are also not accepted
      with self.assertRaisesRegex(_A    ,'foo'      ):
          _lowerCAmelCase       :   List[Any]  =    {'foo': 'bar'}
          model.generate(_A    ,**_A      )








"""simple docstring"""






import inspect
import os
import unittest

import torch

import accelerate
from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
from accelerate.utils import patch_environment







class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    inspect.getfile(accelerate.test_utils      )
      _lowerCAmelCase       :   Union[str, Any]  =    os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ['scripts', 'test_script.py']      )
      _lowerCAmelCase       :   int  =    os.path.sep.join(
          mod_file.split(os.path.sep      )[:-1] + ['scripts', 'test_distributed_data_loop.py']      )
      _lowerCAmelCase       :   Optional[Any]  =    os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ['scripts', 'test_ops.py']      )

  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      print(F"""Found {torch.cuda.device_count()} devices."""      )
      _lowerCAmelCase       :   int  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
      with patch_environment(omp_num_threads=1      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )

  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      print(F"""Found {torch.cuda.device_count()} devices."""      )
      _lowerCAmelCase       :   str  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.operation_file_path]
      print(F"""Command: {cmd}"""      )
      with patch_environment(omp_num_threads=1      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )

  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__      )]
      with patch_environment(omp_num_threads=1      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )







  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      print(F"""Found {torch.cuda.device_count()} devices, using 2 devices only"""      )
      _lowerCAmelCase       :   Tuple  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.data_loop_file_path]
      with patch_environment(omp_num_threads=1    ,cuda_visible_devices='0,1'      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )


if __name__ == "__main__":
 _lowerCAmelCase						=		Accelerator()
 _lowerCAmelCase						=		(accelerator.state.process_index + 2, 1_0)
 _lowerCAmelCase						=		torch.randint(0, 1_0, shape).to(accelerator.device)

 _lowerCAmelCase						=		""""""

 _lowerCAmelCase						=		accelerator.pad_across_processes(tensor)
 if tensora.shape[0] != accelerator.state.num_processes + 1:
  error_msg += F"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
 if not torch.equal(tensora[: accelerator.state.process_index + 2], tensor):
  error_msg += "Tensors have different values."
 if not torch.all(tensora[accelerator.state.process_index + 2 :] == 0):
  error_msg += "Padding was not done with the right value (0)."

 _lowerCAmelCase						=		accelerator.pad_across_processes(tensor, pad_first=True)
 if tensora.shape[0] != accelerator.state.num_processes + 1:
  error_msg += F"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
 _lowerCAmelCase						=		accelerator.state.num_processes - accelerator.state.process_index - 1
 if not torch.equal(tensora[index:], tensor):
  error_msg += "Tensors have different values."
 if not torch.all(tensora[:index] == 0):
  error_msg += "Padding was not done with the right value (0)."

 # Raise error at the end to make sure we don't stop at the first failure.
 if len(error_msg) > 0:
  raise ValueError(error_msg)





