

import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


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

 from transformers.modeling_flax_pytorch_utils import (
     convert_pytorch_state_dict_to_flax,
     load_flax_weights_in_pytorch_model,
 )
 from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
 import torch



@require_flax
class     lowercase						:



 '''simple docstring'''


 def __init__(self       ,       __a       ,       __a=14       ,       __a=7       ,       __a=True       ,       __a=True       ,       __a=True       ,       __a=99       ,       __a=32       ,       __a=5       ,       __a=4       ,       __a=37       ,       __a="gelu"       ,       __a=0.1       ,       __a=0.1       ,       __a=512       ,       __a=0.02       ,       __a=None       ,       ) ->					Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__							=				parent
  UpperCAmelCase__							=				batch_size
  UpperCAmelCase__							=				seq_length
  UpperCAmelCase__							=				is_training
  UpperCAmelCase__							=				use_input_mask
  UpperCAmelCase__							=				use_labels
  UpperCAmelCase__							=				vocab_size
  UpperCAmelCase__							=				d_model
  UpperCAmelCase__							=				num_hidden_layers
  UpperCAmelCase__							=				num_attention_heads
  UpperCAmelCase__							=				ffn_dim
  UpperCAmelCase__							=				activation_function
  UpperCAmelCase__							=				activation_dropout
  UpperCAmelCase__							=				attention_dropout
  UpperCAmelCase__							=				max_position_embeddings
  UpperCAmelCase__							=				initializer_range
  UpperCAmelCase__							=				None
  UpperCAmelCase__							=				0
  UpperCAmelCase__							=				2
  UpperCAmelCase__							=				1




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  UpperCAmelCase__							=				np.clip(ids_tensor([self.batch_size, self.seq_length]       ,       self.vocab_size )       ,       3       ,       self.vocab_size )

  UpperCAmelCase__							=				None
  if self.use_input_mask:
   UpperCAmelCase__							=				random_attention_mask([self.batch_size, self.seq_length] )

  UpperCAmelCase__							=				XGLMConfig(
      vocab_size=self.vocab_size       ,       d_model=self.hidden_size       ,       num_layers=self.num_hidden_layers       ,       attention_heads=self.num_attention_heads       ,       ffn_dim=self.ffn_dim       ,       activation_function=self.activation_function       ,       activation_dropout=self.activation_dropout       ,       attention_dropout=self.attention_dropout       ,       max_position_embeddings=self.max_position_embeddings       ,       initializer_range=self.initializer_range       ,       use_cache=__a       ,       bos_token_id=self.bos_token_id       ,       eos_token_id=self.eos_token_id       ,       pad_token_id=self.pad_token_id       ,       )

  return (config, input_ids, input_mask)




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__							=				self.prepare_config_and_inputs()
  UpperCAmelCase__				,    UpperCAmelCase__				,    UpperCAmelCase__							=				config_and_inputs
  UpperCAmelCase__							=				{'input_ids': input_ids, 'attention_mask': attention_mask}
  return config, inputs_dict




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_config_and_inputs()

  UpperCAmelCase__							=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
  UpperCAmelCase__							=				ids_tensor([self.batch_size, self.seq_length]       ,       vocab_size=2 )

  return (
      config,
      input_ids,
      attention_mask,
      encoder_hidden_states,
      encoder_attention_mask,
  )




 def   UpperCamelCase__							(self       ,       __a       ,       __a       ,       __a       ,       __a ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				20
  UpperCAmelCase__							=				model_class_name(__a )

  UpperCAmelCase__							=				model.init_cache(input_ids.shape[0]       ,       __a )
  UpperCAmelCase__							=				jnp.ones((input_ids.shape[0], max_decoder_length)       ,       dtype='i4' )

  UpperCAmelCase__							=				jnp.broadcast_to(
      jnp.arange(input_ids.shape[-1] - 1 )[None, :]       ,       (input_ids.shape[0], input_ids.shape[-1] - 1) )
  UpperCAmelCase__							=				model(
      input_ids[:, :-1]       ,       attention_mask=__a       ,       past_key_values=__a       ,       position_ids=__a       ,       )

  UpperCAmelCase__							=				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,       dtype='i4' )
  UpperCAmelCase__							=				model(
      input_ids[:, -1:]       ,       attention_mask=__a       ,       past_key_values=outputs_cache.past_key_values       ,       position_ids=__a       ,       )

  UpperCAmelCase__							=				model(__a )

  UpperCAmelCase__							=				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
  self.parent.assertTrue(diff < 1E-3       ,       msg=F"Max diff is {diff}" )




 def   UpperCamelCase__							(self       ,       __a       ,       __a       ,       __a       ,       __a ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__							=				20
  UpperCAmelCase__							=				model_class_name(__a )

  UpperCAmelCase__							=				jnp.concatenate(
      [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1]) )]       ,       axis=-1       ,       )

  UpperCAmelCase__							=				model.init_cache(input_ids.shape[0]       ,       __a )
  UpperCAmelCase__							=				jnp.broadcast_to(
      jnp.arange(input_ids.shape[-1] - 1 )[None, :]       ,       (input_ids.shape[0], input_ids.shape[-1] - 1) )

  UpperCAmelCase__							=				model(
      input_ids[:, :-1]       ,       attention_mask=__a       ,       past_key_values=__a       ,       position_ids=__a       ,       )
  UpperCAmelCase__							=				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,       dtype='i4' )
  UpperCAmelCase__							=				model(
      input_ids[:, -1:]       ,       past_key_values=outputs_cache.past_key_values       ,       attention_mask=__a       ,       position_ids=__a       ,       )

  UpperCAmelCase__							=				model(__a       ,       attention_mask=__a )
  UpperCAmelCase__							=				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
  self.parent.assertTrue(diff < 1E-3       ,       msg=F"Max diff is {diff}" )







@require_sentencepiece
@require_flax
class     lowercase						(       _UpperCamelCase ,				_UpperCamelCase ,				unittest.TestCase					):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE												=				(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
 __SCREAMING_SNAKE_CASE												=				(FlaxXGLMForCausalLM,) if is_flax_available() else ()




 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				FlaxXGLMModelTester(self )




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  for model_class_name in self.all_model_classes:
   UpperCAmelCase__				,    UpperCAmelCase__				,    UpperCAmelCase__							=				self.model_tester.prepare_config_and_inputs()
   self.model_tester.check_use_cache_forward(__a       ,       __a       ,       __a       ,       __a )




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  for model_class_name in self.all_model_classes:
   UpperCAmelCase__				,    UpperCAmelCase__				,    UpperCAmelCase__							=				self.model_tester.prepare_config_and_inputs()
   self.model_tester.check_use_cache_forward_with_attn_mask(
       __a       ,       __a       ,       __a       ,       __a )




 @slow
 def   UpperCamelCase__							(self ) ->					int:






  """simple docstring"""





  UpperCAmelCase__							=				XGLMTokenizer.from_pretrained('XGLM'       ,       padding_side='left' )
  UpperCAmelCase__							=				tokenizer(['Hello this is a long string', 'Hey']       ,       return_tensors='np'       ,       padding=__a       ,       truncation=__a )

  UpperCAmelCase__							=				FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M' )
  UpperCAmelCase__							=				1
  UpperCAmelCase__							=				False

  UpperCAmelCase__							=				jax.jit(model.generate )

  UpperCAmelCase__							=				jit_generate(inputs['input_ids']       ,       attention_mask=inputs['attention_mask'] ).sequences

  UpperCAmelCase__							=				tokenizer.batch_decode(__a       ,       skip_special_tokens=__a )

  UpperCAmelCase__							=				[
      'Hello this is a long string of questions, but I\'m not sure if I\'m',
      'Hey, I\'m a newbie to the forum and I\'',
  ]

  self.assertListEqual(__a       ,       __a )




 @is_pt_flax_cross_test
 def   UpperCamelCase__							(self ) ->					int:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.model_tester.prepare_config_and_inputs_for_common()

  for model_class in self.all_model_classes:
   with self.subTest(model_class.__name__ ):
    # prepare inputs
    UpperCAmelCase__							=				self._prepare_for_class(__a       ,       __a )
    UpperCAmelCase__							=				{k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

    # load corresponding PyTorch class
    UpperCAmelCase__							=				model_class.__name__[4:]  # Skip the "Flax" at the beginning
    UpperCAmelCase__							=				getattr(__a       ,       __a )

    UpperCAmelCase__				,    UpperCAmelCase__							=				pt_inputs['input_ids'].shape
    UpperCAmelCase__							=				np.random.randint(0       ,       seq_length - 1       ,       size=(batch_size,) )
    for batch_idx, start_index in enumerate(__a ):
     UpperCAmelCase__							=				0
     UpperCAmelCase__							=				1
     UpperCAmelCase__							=				0
     UpperCAmelCase__							=				1
    UpperCAmelCase__							=				pt_model_class(__a ).eval()
    # Flax models don't use the `use_cache` option and cache is not returned as a default.
    # So we disable `use_cache` here for PyTorch model.
    UpperCAmelCase__							=				False
    UpperCAmelCase__							=				model_class(__a       ,       dtype=jnp.floataa )
    UpperCAmelCase__							=				convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,       __a )
    UpperCAmelCase__							=				fx_state

    with torch.no_grad():
     UpperCAmelCase__							=				pt_model(**__a ).to_tuple()

    UpperCAmelCase__							=				fx_model(**__a ).to_tuple()
    self.assertEqual(len(__a )       ,       len(__a )       ,       'Output lengths differ between Flax and PyTorch' )
    for fx_output, pt_output in zip(__a       ,       __a ):
     self.assert_almost_equals(fx_output[:, -1]       ,       pt_output[:, -1].numpy()       ,       4E-2 )

    with tempfile.TemporaryDirectory() as tmpdirname:
     pt_model.save_pretrained(__a )
     UpperCAmelCase__							=				model_class.from_pretrained(__a       ,       from_pt=__a )

    UpperCAmelCase__							=				fx_model_loaded(**__a ).to_tuple()
    self.assertEqual(
        len(__a )       ,       len(__a )       ,       'Output lengths differ between Flax and PyTorch' )
    for fx_output_loaded, pt_output in zip(__a       ,       __a ):
     self.assert_almost_equals(fx_output_loaded[:, -1]       ,       pt_output[:, -1].numpy()       ,       4E-2 )




 @is_pt_flax_cross_test
 def   UpperCamelCase__							(self ) ->					str:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.model_tester.prepare_config_and_inputs_for_common()
  for model_class in self.all_model_classes:
   with self.subTest(model_class.__name__ ):
    # prepare inputs
    UpperCAmelCase__							=				self._prepare_for_class(__a       ,       __a )
    UpperCAmelCase__							=				{k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

    # load corresponding PyTorch class
    UpperCAmelCase__							=				model_class.__name__[4:]  # Skip the "Flax" at the beginning
    UpperCAmelCase__							=				getattr(__a       ,       __a )

    UpperCAmelCase__							=				pt_model_class(__a ).eval()
    UpperCAmelCase__							=				False
    UpperCAmelCase__							=				model_class(__a       ,       dtype=jnp.floataa )

    UpperCAmelCase__							=				load_flax_weights_in_pytorch_model(__a       ,       fx_model.params )
    UpperCAmelCase__				,    UpperCAmelCase__							=				pt_inputs['input_ids'].shape
    UpperCAmelCase__							=				np.random.randint(0       ,       seq_length - 1       ,       size=(batch_size,) )
    for batch_idx, start_index in enumerate(__a ):
     UpperCAmelCase__							=				0
     UpperCAmelCase__							=				1
     UpperCAmelCase__							=				0
     UpperCAmelCase__							=				1

    # make sure weights are tied in PyTorch
    pt_model.tie_weights()

    with torch.no_grad():
     UpperCAmelCase__							=				pt_model(**__a ).to_tuple()

    UpperCAmelCase__							=				fx_model(**__a ).to_tuple()
    self.assertEqual(len(__a )       ,       len(__a )       ,       'Output lengths differ between Flax and PyTorch' )
    for fx_output, pt_output in zip(__a       ,       __a ):
     self.assert_almost_equals(fx_output[:, -1]       ,       pt_output[:, -1].numpy()       ,       4E-2 )

    with tempfile.TemporaryDirectory() as tmpdirname:
     fx_model.save_pretrained(__a )
     UpperCAmelCase__							=				pt_model_class.from_pretrained(__a       ,       from_flax=__a )

    with torch.no_grad():
     UpperCAmelCase__							=				pt_model_loaded(**__a ).to_tuple()

    self.assertEqual(
        len(__a )       ,       len(__a )       ,       'Output lengths differ between Flax and PyTorch' )
    for fx_output, pt_output in zip(__a       ,       __a ):
     self.assert_almost_equals(fx_output[:, -1]       ,       pt_output[:, -1].numpy()       ,       4E-2 )




 @slow
 def   UpperCamelCase__							(self ) ->					str:






  """simple docstring"""





  for model_class_name in self.all_model_classes:
   UpperCAmelCase__							=				model_class_name.from_pretrained('facebook/xglm-564M' )
   UpperCAmelCase__							=				model(np.ones((1, 1) ) )
   self.assertIsNotNone(__a )






import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_UpperCamelCase			       =			namedtuple('''ExampleNamedTuple''', '''a b c''')



class     lowercase						(       unittest.TestCase					):



 '''simple docstring'''


 def   UpperCamelCase__							(self ) ->					Tuple:






  """simple docstring"""





  UpperCAmelCase__							=				torch.randn(5       ,       2 )
  UpperCAmelCase__							=				torch.device('cuda' ) if torch.cuda.is_available() else torch.device('cpu' )

  UpperCAmelCase__							=				send_to_device(__a       ,       __a )
  self.assertTrue(torch.equal(resulta.cpu()       ,       __a ) )

  UpperCAmelCase__							=				send_to_device((tensor, [tensor, tensor], 1)       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta[0].cpu()       ,       __a ) )
  self.assertIsInstance(resulta[1]       ,       __a )
  self.assertTrue(torch.equal(resulta[1][0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta[1][1].cpu()       ,       __a ) )
  self.assertEqual(resulta[2]       ,       1 )

  UpperCAmelCase__							=				send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta['a'].cpu()       ,       __a ) )
  self.assertIsInstance(resulta['b']       ,       __a )
  self.assertTrue(torch.equal(resulta['b'][0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta['b'][1].cpu()       ,       __a ) )
  self.assertEqual(resulta['c']       ,       1 )

  UpperCAmelCase__							=				send_to_device(ExampleNamedTuple(a=__a       ,       b=[tensor, tensor]       ,       c=1 )       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta.a.cpu()       ,       __a ) )
  self.assertIsInstance(resulta.b       ,       __a )
  self.assertTrue(torch.equal(resulta.b[0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta.b[1].cpu()       ,       __a ) )
  self.assertEqual(resulta.c       ,       1 )

  UpperCAmelCase__							=				send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1} )       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta['a'].cpu()       ,       __a ) )
  self.assertIsInstance(resulta['b']       ,       __a )
  self.assertTrue(torch.equal(resulta['b'][0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta['b'][1].cpu()       ,       __a ) )
  self.assertEqual(resulta['c']       ,       1 )




 def   UpperCamelCase__							(self ) ->					List[str]:






  """simple docstring"""





  with self.assertRaises(__a ) as cm:
   UpperCAmelCase__							=				recursively_apply(torch.tensor       ,       (torch.tensor(1 ), 1)       ,       error_on_other_type=__a )
  self.assertEqual(
      str(cm.exception )       ,       'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'       ,       )




 def   UpperCamelCase__							(self ) ->					List[str]:






  """simple docstring"""





  UpperCAmelCase__							=				torch.tensor([1, 2, 3, 4, 5] )
  self.assertEqual(listify(__a )       ,       [1, 2, 3, 4, 5] )

  UpperCAmelCase__							=				torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]] )
  self.assertEqual(listify(__a )       ,       [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]] )

  UpperCAmelCase__							=				torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]] )
  self.assertEqual(
      listify(__a )       ,       [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]] )




 def   UpperCamelCase__							(self ) ->					int:






  """simple docstring"""





  with patch_environment(aa=1       ,       BB=2 ):
   self.assertEqual(os.environ.get('AA' )       ,       '1' )
   self.assertEqual(os.environ.get('BB' )       ,       '2' )

  self.assertNotIn('AA'       ,       os.environ )
  self.assertNotIn('BB'       ,       os.environ )




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				model.forward
  UpperCAmelCase__							=				convert_outputs_to_fpaa(model.forward )
  UpperCAmelCase__							=				extract_model_from_parallel(__a       ,       keep_fpaa_wrapper=__a )
  UpperCAmelCase__							=				pickle.dumps(__a )




 @require_cuda
 def   UpperCamelCase__							(self ) ->					Tuple:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				model.forward
  UpperCAmelCase__							=				torch.cuda.amp.autocast(dtype=torch.floataa )(model.forward )
  UpperCAmelCase__							=				convert_outputs_to_fpaa(model.forward )
  UpperCAmelCase__							=				extract_model_from_parallel(__a       ,       keep_fpaa_wrapper=__a )
  UpperCAmelCase__							=				pickle.dumps(__a )




 @require_cuda
 @require_torch_min_version(version='2.0' )
 def   UpperCamelCase__							(self ) ->					List[str]:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				model.forward
  UpperCAmelCase__							=				torch.cuda.amp.autocast(dtype=torch.floataa )(model.forward )
  UpperCAmelCase__							=				convert_outputs_to_fpaa(model.forward )
  UpperCAmelCase__							=				torch.compile(model.forward       ,       backend='inductor' )
  UpperCAmelCase__							=				torch.randn(4       ,       10 ).cuda()
  UpperCAmelCase__							=				model(__a )




 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
  UpperCAmelCase__							=				torch.nn.parallel.DataParallel(__a )
  UpperCAmelCase__							=				extract_model_from_parallel(__a )

  self.assertEqual(__a       ,       __a )




 @require_torch_min_version(version='2.0' )
 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				torch.compile(__a )

  # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
  UpperCAmelCase__							=				torch.nn.parallel.DataParallel(__a )
  UpperCAmelCase__							=				torch.compile(__a )
  UpperCAmelCase__							=				extract_model_from_parallel(__a )

  self.assertEqual(compiled_model._orig_mod       ,       compiled_model_unwrapped._orig_mod )




 def   UpperCamelCase__							(self ) ->					str:






  """simple docstring"""





  self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3] )] )       ,       torch.device('cpu' ) )
  self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3] )} )       ,       torch.device('cpu' ) )
  self.assertIsNone(find_device([1, 'a'] ) )



