




"""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							:			Tuple     							=							"""0.12"""  # assumed parallelism: 8


if is_torch_available():
					import torch



def     a__ (				snake_case__       , snake_case__       , snake_case__=None				)						->					Dict:
			if rng is None:
						lowerCamelCase			      =					random.Random()

			lowerCamelCase			      =					1
			for dim in shape:
						total_dims *= dim

			lowerCamelCase			      =					[]
			for _ in range(snake_case__				):
						values.append(rng.randint(0       , vocab_size - 1				)				)

			lowerCamelCase			      =					np.array(snake_case__       , dtype=jnp.intaa				).reshape(snake_case__				)

			return output



def     a__ (				snake_case__       , snake_case__=None				)						->					int:
			lowerCamelCase			      =					ids_tensor(snake_case__       , vocab_size=2       , rng=snake_case__				)
			# make sure that at least one token is attended to for each batch
			lowerCamelCase			      =					1
			return attn_mask






@require_flax
class     __magic_name__							:






				'''simple docstring'''
				__UpperCamelCase			  =						None
				__UpperCamelCase			  =						()





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs_for_common()

							# cut to half length & take max batch_size 3
							lowerCamelCase			      =					2
							lowerCamelCase			      =					inputs["""input_ids"""].shape[-1] // 2
							lowerCamelCase			      =					inputs["""input_ids"""][:max_batch_size, :sequence_length]

							lowerCamelCase			      =					jnp.ones_like(_a )
							lowerCamelCase			      =					attention_mask[:max_batch_size, :sequence_length]

							# generate max 5 tokens
							lowerCamelCase			      =					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			      =					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			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					False
							lowerCamelCase			      =					max_length
							lowerCamelCase			      =					0

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model_class.__name__[4:]  # Skip the "Flax" at the beginning
										lowerCamelCase			      =					getattr(_a		, _a )
										lowerCamelCase			      =					pt_model_class(_a ).eval()
										lowerCamelCase			      =					load_flax_weights_in_pytorch_model(_a		, flax_model.params )

										lowerCamelCase			      =					flax_model.generate(_a ).sequences
										lowerCamelCase			      =					pt_model.generate(torch.tensor(_a		, dtype=torch.long ) )

										if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
													lowerCamelCase			      =					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			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					False
							lowerCamelCase			      =					max_length

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a ).sequences

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





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					True
							lowerCamelCase			      =					max_length

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a ).sequences

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





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					False
							lowerCamelCase			      =					max_length
							lowerCamelCase			      =					2

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a ).sequences

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





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					False
							lowerCamelCase			      =					max_length
							lowerCamelCase			      =					2
							lowerCamelCase			      =					2

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					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			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					True
							lowerCamelCase			      =					max_length
							lowerCamelCase			      =					0.8
							lowerCamelCase			      =					10
							lowerCamelCase			      =					0.3
							lowerCamelCase			      =					1
							lowerCamelCase			      =					8
							lowerCamelCase			      =					9

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a ).sequences

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





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					max_length
							lowerCamelCase			      =					1
							lowerCamelCase			      =					8
							lowerCamelCase			      =					9

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a ).sequences

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





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					self._get_input_ids_and_config()
							lowerCamelCase			      =					max_length
							lowerCamelCase			      =					2
							lowerCamelCase			      =					1
							lowerCamelCase			      =					8
							lowerCamelCase			      =					9

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a ).sequences

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





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					self._get_input_ids_and_config()

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

							lowerCamelCase			      =					False
							lowerCamelCase			      =					max_length

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a		, attention_mask=_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					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			      =					self._get_input_ids_and_config()

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

							lowerCamelCase			      =					True
							lowerCamelCase			      =					max_length

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a		, attention_mask=_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					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			      =					self._get_input_ids_and_config()

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

							lowerCamelCase			      =					2
							lowerCamelCase			      =					max_length

							for model_class in self.all_generative_model_classes:
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model.generate(_a		, attention_mask=_a ).sequences
										self.assertEqual(generation_outputs.shape[-1]		, _a )

										lowerCamelCase			      =					jit(model.generate )
										lowerCamelCase			      =					jit_generate(_a		, attention_mask=_a ).sequences

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






@require_flax
class     __magic_name__							(       unittest.TestCase				):






				'''simple docstring'''





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert""" )
							lowerCamelCase			      =					FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only""" )

							lowerCamelCase			      =					"""Hello world"""
							lowerCamelCase			      =					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			      =					{"""foo""": """bar"""}
										model.generate(_a		, **_a )









"""simple docstring"""




import operator as op


lowerCAmelCase							:			Dict     							=							"""scaler.pt"""
lowerCAmelCase							:			Tuple     							=							"""pytorch_model"""
lowerCAmelCase							:			Union[str, Any]     							=							"""random_states"""
lowerCAmelCase							:			Union[str, Any]     							=							"""optimizer"""
lowerCAmelCase							:			Dict     							=							"""scheduler"""
lowerCAmelCase							:			int     							=							"""pytorch_model.bin"""
lowerCAmelCase							:			str     							=							"""pytorch_model.bin.index.json"""
lowerCAmelCase							:			Union[str, Any]     							=							"""model.safetensors"""
lowerCAmelCase							:			List[Any]     							=							"""model.safetensors.index.json"""
lowerCAmelCase							:			List[Any]     							=							"""1.10.2"""
lowerCAmelCase							:			Any     							=							"""py38"""
lowerCAmelCase							:			Optional[int]     							=							"""4.17.0"""
lowerCAmelCase							:			str     							=							["""ml.p3.16xlarge""", """ml.p3dn.24xlarge""", """ml.p4dn.24xlarge"""]
lowerCAmelCase							:			Tuple     							=							["""FULL_SHARD""", """SHARD_GRAD_OP""", """NO_SHARD""", """HYBRID_SHARD""", """HYBRID_SHARD_ZERO2"""]
lowerCAmelCase							:			List[Any]     							=							["""TRANSFORMER_BASED_WRAP""", """SIZE_BASED_WRAP""", """NO_WRAP"""]
lowerCAmelCase							:			List[str]     							=							["""BACKWARD_PRE""", """BACKWARD_POST""", """NO_PREFETCH"""]
lowerCAmelCase							:			List[str]     							=							["""FULL_STATE_DICT""", """LOCAL_STATE_DICT""", """SHARDED_STATE_DICT"""]
lowerCAmelCase							:			Any     							=							"""2.0.1"""
lowerCAmelCase							:			List[Any]     							=							["""pdsh""", """standard""", """openmpi""", """mvapich"""]
lowerCAmelCase							:			Union[str, Any]     							=							["""default""", """reduce-overhead""", """max-autotune"""]

lowerCAmelCase							:			Optional[int]     							=							{""">""": op.gt, """>=""": op.ge, """==""": op.eq, """!=""": op.ne, """<=""": op.le, """<""": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
lowerCAmelCase							:			Union[str, Any]     							=							[
    """nnodes""",
    """nproc_per_node""",
    """rdzv_backend""",
    """rdzv_endpoint""",
    """rdzv_id""",
    """rdzv_conf""",
    """standalone""",
    """max_restarts""",
    """monitor_interval""",
    """start_method""",
    """role""",
    """module""",
    """m""",
    """no_python""",
    """run_path""",
    """log_dir""",
    """r""",
    """redirects""",
    """t""",
    """tee""",
    """node_rank""",
    """master_addr""",
    """master_port""",
]

lowerCAmelCase							:			List[str]     							=							["""DEEPSPEED""", """MULTI_GPU""", """FSDP""", """MEGATRON_LM"""]
lowerCAmelCase							:			Optional[Any]     							=							["""DEEPSPEED""", """MULTI_XPU""", """FSDP"""]



