code: 
  
  
  
  
  
  
  import tempfile
  import unittest
  
  import numpy as np
  
  import transformers
  from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
  from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow
  
  from ...generation.test_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask
  
  
  if is_flax_available():
     import jax
     import jax.numpy as jnp
  
     from transformers.modeling_flax_pytorch_utils import (
         convert_pytorch_state_dict_to_flax,
         load_flax_weights_in_pytorch_model,
     )
     from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel
  
  if is_torch_available():
     import torch
  
  class 					snake_case_							:
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
     def __init__(    self:	Any   ,       snake_case__:	List[str]   ,       snake_case__:	str=14   ,       snake_case__:	int=7   ,       snake_case__:	Dict=True   ,       snake_case__:	Optional[Any]=True   ,       snake_case__:	Optional[int]=False   ,       snake_case__:	str=True   ,       snake_case__:	str=99   ,       snake_case__:	Union[str, Any]=32   ,       snake_case__:	Union[str, Any]=4   ,       snake_case__:	str=4   ,       snake_case__:	Dict=[[["global", "local"], 2]]   ,       snake_case__:	List[Any]=37   ,       snake_case__:	Optional[int]="gelu"   ,       snake_case__:	Dict=0.1   ,       snake_case__:	int=0.1   ,       snake_case__:	str=512   ,       snake_case__:	Any=7   ,       snake_case__:	Dict=0.02   ,       )  ->		Optional[int]:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	parent
         _lowercase     =	batch_size
         _lowercase     =	seq_length
         _lowercase     =	is_training
         _lowercase     =	use_input_mask
         _lowercase     =	use_token_type_ids
         _lowercase     =	use_labels
         _lowercase     =	vocab_size
         _lowercase     =	hidden_size
         _lowercase     =	num_hidden_layers
         _lowercase     =	num_attention_heads
         _lowercase     =	attention_types
         _lowercase     =	intermediate_size
         _lowercase     =	hidden_act
         _lowercase     =	hidden_dropout_prob
         _lowercase     =	attention_probs_dropout_prob
         _lowercase     =	max_position_embeddings
         _lowercase     =	window_size
         _lowercase     =	initializer_range
         _lowercase     =	None
         _lowercase     =	vocab_size - 1
         _lowercase     =	vocab_size - 1
         _lowercase     =	vocab_size - 1
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	List[Any]	)  ->		List[Any]:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	ids_tensor([self.batch_size, self.seq_length]   ,       self.vocab_size	)
  
         _lowercase     =	None
         if self.use_input_mask:
             _lowercase     =	random_attention_mask([self.batch_size, self.seq_length]	)
  
         _lowercase     =	GPTNeoConfig(
             vocab_size=self.vocab_size   ,       hidden_size=self.hidden_size   ,       num_layers=self.num_hidden_layers   ,       num_heads=self.num_attention_heads   ,       max_position_embeddings=self.max_position_embeddings   ,       use_cache=snake_case__   ,       bos_token_id=self.bos_token_id   ,       eos_token_id=self.eos_token_id   ,       pad_token_id=self.pad_token_id   ,       window_size=self.window_size   ,       attention_types=self.attention_types   ,       )
  
         return (config, input_ids, input_mask)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	str	)  ->		Dict:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	self.prepare_config_and_inputs()
         _lowercase ,    _lowercase ,    _lowercase     =	config_and_inputs
         _lowercase     =	{'input_ids': input_ids, 'attention_mask': attention_mask}
         return config, inputs_dict
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Union[str, Any]   ,       snake_case__:	List[Any]   ,       snake_case__:	List[str]   ,       snake_case__:	Tuple   ,       snake_case__:	Optional[int]	)  ->		Optional[int]:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	20
         _lowercase     =	model_class_name(snake_case__	)
  
         _lowercase     =	model.init_cache(input_ids.shape[0]   ,       snake_case__	)
         _lowercase     =	jnp.ones((input_ids.shape[0], max_decoder_length)   ,       dtype='i4'	)
  
         _lowercase     =	jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1	)[None, :]   ,       (input_ids.shape[0], input_ids.shape[-1] - 1)	)
         _lowercase     =	model(
             input_ids[:, :-1]   ,       attention_mask=snake_case__   ,       past_key_values=snake_case__   ,       position_ids=snake_case__   ,       )
  
         _lowercase     =	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]   ,       dtype='i4'	)
         _lowercase     =	model(
             input_ids[:, -1:]   ,       attention_mask=snake_case__   ,       past_key_values=outputs_cache.past_key_values   ,       position_ids=snake_case__   ,       )
  
         _lowercase     =	model(snake_case__	)
  
         _lowercase     =	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 			SCREAMING_SNAKE_CASE__				(    self:	Tuple   ,       snake_case__:	List[Any]   ,       snake_case__:	int   ,       snake_case__:	Tuple   ,       snake_case__:	Tuple	)  ->		int:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	20
         _lowercase     =	model_class_name(snake_case__	)
  
         _lowercase     =	jnp.concatenate(
             [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)]   ,       axis=-1   ,       )
  
         _lowercase     =	model.init_cache(input_ids.shape[0]   ,       snake_case__	)
         _lowercase     =	jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1	)[None, :]   ,       (input_ids.shape[0], input_ids.shape[-1] - 1)	)
  
         _lowercase     =	model(
             input_ids[:, :-1]   ,       attention_mask=snake_case__   ,       past_key_values=snake_case__   ,       position_ids=snake_case__   ,       )
         _lowercase     =	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]   ,       dtype='i4'	)
         _lowercase     =	model(
             input_ids[:, -1:]   ,       past_key_values=outputs_cache.past_key_values   ,       attention_mask=snake_case__   ,       position_ids=snake_case__   ,       )
  
         _lowercase     =	model(snake_case__   ,       attention_mask=snake_case__	)
  
         _lowercase     =	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_flax
  class 					snake_case_							(       lowercase_      ,			lowercase_      ,			unittest.TestCase      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
     __lowerCAmelCase      =						(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
     __lowerCAmelCase      =						(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	List[str]	)  ->		Any:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	FlaxGPTNeoModelTester(self	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Union[str, Any]	)  ->		Tuple:
  
  
         """simple docstring"""
  
  
  
         for model_class_name in self.all_model_classes:
             _lowercase ,    _lowercase ,    _lowercase     =	self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward(snake_case__   ,       snake_case__   ,       snake_case__   ,       snake_case__	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Tuple	)  ->		Any:
  
  
         """simple docstring"""
  
  
  
         for model_class_name in self.all_model_classes:
             _lowercase ,    _lowercase ,    _lowercase     =	self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward_with_attn_mask(
                 snake_case__   ,       snake_case__   ,       snake_case__   ,       snake_case__	)
  
     @slow
     def 			SCREAMING_SNAKE_CASE__				(    self:	str	)  ->		List[str]:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	GPTaTokenizer.from_pretrained('gpt2'   ,       pad_token='<|endoftext|>'   ,       padding_side='left'	)
         _lowercase     =	tokenizer(['Hello this is a long string', 'Hey']   ,       return_tensors='np'   ,       padding=snake_case__   ,       truncation=snake_case__	)
  
         _lowercase     =	FlaxGPTNeoForCausalLM.from_pretrained('EleutherAI/gpt-neo-125M'	)
         _lowercase     =	False
         _lowercase     =	model.config.eos_token_id
  
         _lowercase     =	jax.jit(model.generate	)
  
         _lowercase     =	jit_generate(
             inputs['input_ids']   ,       attention_mask=inputs['attention_mask']   ,       pad_token_id=tokenizer.pad_token_id	).sequences
  
         _lowercase     =	tokenizer.batch_decode(snake_case__   ,       skip_special_tokens=snake_case__	)
  
         _lowercase     =	[
             'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
             'Hey, I\'m a little late to the party. I\'m going to',
         ]
  
         self.assertListEqual(snake_case__   ,       snake_case__	)
  
     @is_pt_flax_cross_test
     def 			SCREAMING_SNAKE_CASE__				(    self:	Any	)  ->		Union[str, Any]:
  
  
         """simple docstring"""
  
  
  
         _lowercase ,    _lowercase     =	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
                 _lowercase     =	self._prepare_for_class(snake_case__   ,       snake_case__	)
                 _lowercase     =	{k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}
  
                 # load corresponding PyTorch class
                 _lowercase     =	model_class.__name__[4:]  # Skip the "Flax" at the beginning
                 _lowercase     =	getattr(snake_case__   ,       snake_case__	)
  
                 _lowercase ,    _lowercase     =	pt_inputs['input_ids'].shape
                 _lowercase     =	np.random.randint(0   ,       seq_length - 1   ,       size=(batch_size,)	)
                 for batch_idx, start_index in enumerate(snake_case__	):
                     _lowercase     =	0
                     _lowercase     =	1
                     _lowercase     =	0
                     _lowercase     =	1
                 _lowercase     =	pt_model_class(snake_case__	).eval()
                 _lowercase     =	model_class(snake_case__   ,       dtype=jnp.floataa	)
  
                 _lowercase     =	convert_pytorch_state_dict_to_flax(pt_model.state_dict()   ,       snake_case__	)
                 _lowercase     =	fx_state
  
                 with torch.no_grad():
                     _lowercase     =	pt_model(**snake_case__	).to_tuple()
  
                 _lowercase     =	fx_model(**snake_case__	).to_tuple()
                 self.assertEqual(len(snake_case__	)   ,       len(snake_case__	)   ,       'Output lengths differ between Flax and PyTorch'	)
                 for fx_output, pt_output in zip(snake_case__   ,       snake_case__	):
                     self.assert_almost_equals(fx_output[:, -1]   ,       pt_output[:, -1].numpy()   ,       4E-2	)
  
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     pt_model.save_pretrained(snake_case__	)
                     _lowercase     =	model_class.from_pretrained(snake_case__   ,       from_pt=snake_case__	)
  
                 _lowercase     =	fx_model_loaded(**snake_case__	).to_tuple()
                 self.assertEqual(
                     len(snake_case__	)   ,       len(snake_case__	)   ,       'Output lengths differ between Flax and PyTorch'	)
                 for fx_output_loaded, pt_output in zip(snake_case__   ,       snake_case__	):
                     self.assert_almost_equals(fx_output_loaded[:, -1]   ,       pt_output[:, -1].numpy()   ,       4E-2	)
  
     @is_pt_flax_cross_test
     def 			SCREAMING_SNAKE_CASE__				(    self:	List[str]	)  ->		Optional[int]:
  
  
         """simple docstring"""
  
  
  
         _lowercase ,    _lowercase     =	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
                 _lowercase     =	self._prepare_for_class(snake_case__   ,       snake_case__	)
                 _lowercase     =	{k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}
  
                 # load corresponding PyTorch class
                 _lowercase     =	model_class.__name__[4:]  # Skip the "Flax" at the beginning
                 _lowercase     =	getattr(snake_case__   ,       snake_case__	)
  
                 _lowercase     =	pt_model_class(snake_case__	).eval()
                 _lowercase     =	model_class(snake_case__   ,       dtype=jnp.floataa	)
  
                 _lowercase     =	load_flax_weights_in_pytorch_model(snake_case__   ,       fx_model.params	)
                 _lowercase ,    _lowercase     =	pt_inputs['input_ids'].shape
                 _lowercase     =	np.random.randint(0   ,       seq_length - 1   ,       size=(batch_size,)	)
                 for batch_idx, start_index in enumerate(snake_case__	):
                     _lowercase     =	0
                     _lowercase     =	1
                     _lowercase     =	0
                     _lowercase     =	1
  
                 # make sure weights are tied in PyTorch
                 pt_model.tie_weights()
  
                 with torch.no_grad():
                     _lowercase     =	pt_model(**snake_case__	).to_tuple()
  
                 _lowercase     =	fx_model(**snake_case__	).to_tuple()
                 self.assertEqual(len(snake_case__	)   ,       len(snake_case__	)   ,       'Output lengths differ between Flax and PyTorch'	)
                 for fx_output, pt_output in zip(snake_case__   ,       snake_case__	):
                     self.assert_almost_equals(fx_output[:, -1]   ,       pt_output[:, -1].numpy()   ,       4E-2	)
  
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     fx_model.save_pretrained(snake_case__	)
                     _lowercase     =	pt_model_class.from_pretrained(snake_case__   ,       from_flax=snake_case__	)
  
                 with torch.no_grad():
                     _lowercase     =	pt_model_loaded(**snake_case__	).to_tuple()
  
                 self.assertEqual(
                     len(snake_case__	)   ,       len(snake_case__	)   ,       'Output lengths differ between Flax and PyTorch'	)
                 for fx_output, pt_output in zip(snake_case__   ,       snake_case__	):
                     self.assert_almost_equals(fx_output[:, -1]   ,       pt_output[:, -1].numpy()   ,       4E-2	)
  
     @slow
     def 			SCREAMING_SNAKE_CASE__				(    self:	Optional[Any]	)  ->		Dict:
  
  
         """simple docstring"""
  
  
  
         for model_class_name in self.all_model_classes:
             _lowercase     =	model_class_name.from_pretrained('EleutherAI/gpt-neo-125M'	)
             _lowercase     =	model(np.ones((1, 1)	)	)
             self.assertIsNotNone(snake_case__	)
  
style_context: 
  
  
  
  
  
  
  from __future__ import annotations
  def        UpperCAmelCase      (   _lowercase,      _lowercase      ):
  
  
      '''simple docstring'''
  
  
  
      _lowercase     =	[None] * (len(_lowercase      ) + len(_lowercase      ))
  
      _lowercase     =	0  # pointer to current index for left Half
      _lowercase     =	0  # pointer to current index for the right Half
      _lowercase     =	0  # pointer to current index for the sorted array Half
  
      while pointera < len(_lowercase      ) and pointera < len(_lowercase      ):
          if left_half[pointera] < right_half[pointera]:
              _lowercase     =	left_half[pointera]
              pointera += 1
              index += 1
          else:
              _lowercase     =	right_half[pointera]
              pointera += 1
              index += 1
      while pointera < len(_lowercase      ):
          _lowercase     =	left_half[pointera]
          pointera += 1
          index += 1
  
      while pointera < len(_lowercase      ):
          _lowercase     =	right_half[pointera]
          pointera += 1
          index += 1
  
      return sorted_array
  def        UpperCAmelCase      (   _lowercase      ):
  
  
      '''simple docstring'''
  
  
  
      if len(_lowercase      ) <= 1:
          return array
      # the actual formula to calculate the middle element = left + (right - left) // 2
      # this avoids integer overflow in case of large N
      _lowercase     =	0 + (len(_lowercase      ) - 0) // 2
  
      # Split the array into halves till the array length becomes equal to One
      # merge the arrays of single length returned by mergeSort function and
      # pass them into the merge arrays function which merges the array
      _lowercase     =	array[:middle]
      _lowercase     =	array[middle:]
  
      return merge(merge_sort(_lowercase      ),      merge_sort(_lowercase      )      )
  
  
  if __name__ == "__main__":
     import doctest
  
     doctest.testmod()
  