



'''simple docstring'''




from binascii import hexlify
from hashlib import shaaaa
from os import urandom

# RFC 3526 - More Modular Exponential (MODP) Diffie-Hellman groups for
# Internet Key Exchange (IKE) https://tools.ietf.org/html/rfc3526

a   :			List[Any]					  =    {
    # 1536-bit
    5: {
        """prime""": int(
            """FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"""
            + """29024E088A67CC74020BBEA63B139B22514A08798E3404DD"""
            + """EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"""
            + """E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"""
            + """EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"""
            + """C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"""
            + """83655D23DCA3AD961C62F356208552BB9ED529077096966D"""
            + """670C354E4ABC9804F1746C08CA237327FFFFFFFFFFFFFFFF""",
            base=1_6,
        ),
        """generator""": 2,
    },
    # 2048-bit
    1_4: {
        """prime""": int(
            """FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"""
            + """29024E088A67CC74020BBEA63B139B22514A08798E3404DD"""
            + """EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"""
            + """E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"""
            + """EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"""
            + """C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"""
            + """83655D23DCA3AD961C62F356208552BB9ED529077096966D"""
            + """670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"""
            + """E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"""
            + """DE2BCBF6955817183995497CEA956AE515D2261898FA0510"""
            + """15728E5A8AACAA68FFFFFFFFFFFFFFFF""",
            base=1_6,
        ),
        """generator""": 2,
    },
    # 3072-bit
    1_5: {
        """prime""": int(
            """FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"""
            + """29024E088A67CC74020BBEA63B139B22514A08798E3404DD"""
            + """EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"""
            + """E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"""
            + """EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"""
            + """C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"""
            + """83655D23DCA3AD961C62F356208552BB9ED529077096966D"""
            + """670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"""
            + """E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"""
            + """DE2BCBF6955817183995497CEA956AE515D2261898FA0510"""
            + """15728E5A8AAAC42DAD33170D04507A33A85521ABDF1CBA64"""
            + """ECFB850458DBEF0A8AEA71575D060C7DB3970F85A6E1E4C7"""
            + """ABF5AE8CDB0933D71E8C94E04A25619DCEE3D2261AD2EE6B"""
            + """F12FFA06D98A0864D87602733EC86A64521F2B18177B200C"""
            + """BBE117577A615D6C770988C0BAD946E208E24FA074E5AB31"""
            + """43DB5BFCE0FD108E4B82D120A93AD2CAFFFFFFFFFFFFFFFF""",
            base=1_6,
        ),
        """generator""": 2,
    },
    # 4096-bit
    1_6: {
        """prime""": int(
            """FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"""
            + """29024E088A67CC74020BBEA63B139B22514A08798E3404DD"""
            + """EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"""
            + """E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"""
            + """EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"""
            + """C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"""
            + """83655D23DCA3AD961C62F356208552BB9ED529077096966D"""
            + """670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"""
            + """E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"""
            + """DE2BCBF6955817183995497CEA956AE515D2261898FA0510"""
            + """15728E5A8AAAC42DAD33170D04507A33A85521ABDF1CBA64"""
            + """ECFB850458DBEF0A8AEA71575D060C7DB3970F85A6E1E4C7"""
            + """ABF5AE8CDB0933D71E8C94E04A25619DCEE3D2261AD2EE6B"""
            + """F12FFA06D98A0864D87602733EC86A64521F2B18177B200C"""
            + """BBE117577A615D6C770988C0BAD946E208E24FA074E5AB31"""
            + """43DB5BFCE0FD108E4B82D120A92108011A723C12A787E6D7"""
            + """88719A10BDBA5B2699C327186AF4E23C1A946834B6150BDA"""
            + """2583E9CA2AD44CE8DBBBC2DB04DE8EF92E8EFC141FBECAA6"""
            + """287C59474E6BC05D99B2964FA090C3A2233BA186515BE7ED"""
            + """1F612970CEE2D7AFB81BDD762170481CD0069127D5B05AA9"""
            + """93B4EA988D8FDDC186FFB7DC90A6C08F4DF435C934063199"""
            + """FFFFFFFFFFFFFFFF""",
            base=1_6,
        ),
        """generator""": 2,
    },
    # 6144-bit
    1_7: {
        """prime""": int(
            """FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E08"""
            + """8A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B"""
            + """302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9"""
            + """A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE6"""
            + """49286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8"""
            + """FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D"""
            + """670C354E4ABC9804F1746C08CA18217C32905E462E36CE3BE39E772C"""
            + """180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF695581718"""
            + """3995497CEA956AE515D2261898FA051015728E5A8AAAC42DAD33170D"""
            + """04507A33A85521ABDF1CBA64ECFB850458DBEF0A8AEA71575D060C7D"""
            + """B3970F85A6E1E4C7ABF5AE8CDB0933D71E8C94E04A25619DCEE3D226"""
            + """1AD2EE6BF12FFA06D98A0864D87602733EC86A64521F2B18177B200C"""
            + """BBE117577A615D6C770988C0BAD946E208E24FA074E5AB3143DB5BFC"""
            + """E0FD108E4B82D120A92108011A723C12A787E6D788719A10BDBA5B26"""
            + """99C327186AF4E23C1A946834B6150BDA2583E9CA2AD44CE8DBBBC2DB"""
            + """04DE8EF92E8EFC141FBECAA6287C59474E6BC05D99B2964FA090C3A2"""
            + """233BA186515BE7ED1F612970CEE2D7AFB81BDD762170481CD0069127"""
            + """D5B05AA993B4EA988D8FDDC186FFB7DC90A6C08F4DF435C934028492"""
            + """36C3FAB4D27C7026C1D4DCB2602646DEC9751E763DBA37BDF8FF9406"""
            + """AD9E530EE5DB382F413001AEB06A53ED9027D831179727B0865A8918"""
            + """DA3EDBEBCF9B14ED44CE6CBACED4BB1BDB7F1447E6CC254B33205151"""
            + """2BD7AF426FB8F401378CD2BF5983CA01C64B92ECF032EA15D1721D03"""
            + """F482D7CE6E74FEF6D55E702F46980C82B5A84031900B1C9E59E7C97F"""
            + """BEC7E8F323A97A7E36CC88BE0F1D45B7FF585AC54BD407B22B4154AA"""
            + """CC8F6D7EBF48E1D814CC5ED20F8037E0A79715EEF29BE32806A1D58B"""
            + """B7C5DA76F550AA3D8A1FBFF0EB19CCB1A313D55CDA56C9EC2EF29632"""
            + """387FE8D76E3C0468043E8F663F4860EE12BF2D5B0B7474D6E694F91E"""
            + """6DCC4024FFFFFFFFFFFFFFFF""",
            base=1_6,
        ),
        """generator""": 2,
    },
    # 8192-bit
    1_8: {
        """prime""": int(
            """FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"""
            + """29024E088A67CC74020BBEA63B139B22514A08798E3404DD"""
            + """EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"""
            + """E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"""
            + """EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"""
            + """C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"""
            + """83655D23DCA3AD961C62F356208552BB9ED529077096966D"""
            + """670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"""
            + """E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"""
            + """DE2BCBF6955817183995497CEA956AE515D2261898FA0510"""
            + """15728E5A8AAAC42DAD33170D04507A33A85521ABDF1CBA64"""
            + """ECFB850458DBEF0A8AEA71575D060C7DB3970F85A6E1E4C7"""
            + """ABF5AE8CDB0933D71E8C94E04A25619DCEE3D2261AD2EE6B"""
            + """F12FFA06D98A0864D87602733EC86A64521F2B18177B200C"""
            + """BBE117577A615D6C770988C0BAD946E208E24FA074E5AB31"""
            + """43DB5BFCE0FD108E4B82D120A92108011A723C12A787E6D7"""
            + """88719A10BDBA5B2699C327186AF4E23C1A946834B6150BDA"""
            + """2583E9CA2AD44CE8DBBBC2DB04DE8EF92E8EFC141FBECAA6"""
            + """287C59474E6BC05D99B2964FA090C3A2233BA186515BE7ED"""
            + """1F612970CEE2D7AFB81BDD762170481CD0069127D5B05AA9"""
            + """93B4EA988D8FDDC186FFB7DC90A6C08F4DF435C934028492"""
            + """36C3FAB4D27C7026C1D4DCB2602646DEC9751E763DBA37BD"""
            + """F8FF9406AD9E530EE5DB382F413001AEB06A53ED9027D831"""
            + """179727B0865A8918DA3EDBEBCF9B14ED44CE6CBACED4BB1B"""
            + """DB7F1447E6CC254B332051512BD7AF426FB8F401378CD2BF"""
            + """5983CA01C64B92ECF032EA15D1721D03F482D7CE6E74FEF6"""
            + """D55E702F46980C82B5A84031900B1C9E59E7C97FBEC7E8F3"""
            + """23A97A7E36CC88BE0F1D45B7FF585AC54BD407B22B4154AA"""
            + """CC8F6D7EBF48E1D814CC5ED20F8037E0A79715EEF29BE328"""
            + """06A1D58BB7C5DA76F550AA3D8A1FBFF0EB19CCB1A313D55C"""
            + """DA56C9EC2EF29632387FE8D76E3C0468043E8F663F4860EE"""
            + """12BF2D5B0B7474D6E694F91E6DBE115974A3926F12FEE5E4"""
            + """38777CB6A932DF8CD8BEC4D073B931BA3BC832B68D9DD300"""
            + """741FA7BF8AFC47ED2576F6936BA424663AAB639C5AE4F568"""
            + """3423B4742BF1C978238F16CBE39D652DE3FDB8BEFC848AD9"""
            + """22222E04A4037C0713EB57A81A23F0C73473FC646CEA306B"""
            + """4BCBC8862F8385DDFA9D4B7FA2C087E879683303ED5BDD3A"""
            + """062B3CF5B3A278A66D2A13F83F44F82DDF310EE074AB6A36"""
            + """4597E899A0255DC164F31CC50846851DF9AB48195DED7EA1"""
            + """B1D510BD7EE74D73FAF36BC31ECFA268359046F4EB879F92"""
            + """4009438B481C6CD7889A002ED5EE382BC9190DA6FC026E47"""
            + """9558E4475677E9AA9E3050E2765694DFC81F56E880B96E71"""
            + """60C980DD98EDD3DFFFFFFFFFFFFFFFFF""",
            base=1_6,
        ),
        """generator""": 2,
    },
}

class 			UpperCamelCase_ :





	def __init__(  self		,		A = 14    )			->			None:
								if group not in primes:
															raise ValueError("""Unsupported Group"""    )
								UpperCAmelCase     :		Union[str, Any]							     =		primes[group]["""prime"""]
								UpperCAmelCase     :		Any							     =		primes[group]["""generator"""]

								UpperCAmelCase     :		Optional[Any]							     =		int(hexlify(urandom(32    )    )		,		base=16    )





	def  _lowercase(  self    )			->			str:
								return hex(self.__private_key    )[2:]





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     :		Optional[int]							     =		pow(self.generator		,		self.__private_key		,		self.prime    )
								return hex(A    )[2:]





	def  _lowercase(  self		,		A    )			->			bool:
								# check if the other public key is valid based on NIST SP800-56
								return (
								    2 <= key <= self.prime - 2
								    and pow(A		,		(self.prime - 1) // 2		,		self.prime    ) == 1
								)





	def  _lowercase(  self		,		A    )			->			str:
								UpperCAmelCase     :		Tuple							     =		int(A		,		base=16    )
								if not self.is_valid_public_key(A    ):
															raise ValueError("""Invalid public key"""    )
								UpperCAmelCase     :		Any							     =		pow(A		,		self.__private_key		,		self.prime    )
								return shaaaa(str(A    ).encode()    ).hexdigest()





	@staticmethod
	def  _lowercase(  A		,		A    )			->			bool:
								# check if the other public key is valid based on NIST SP800-56
								return (
								    2 <= remote_public_key_str <= prime - 2
								    and pow(A		,		(prime - 1) // 2		,		A    ) == 1
								)






	@staticmethod
	def  _lowercase(  A		,		A		,		A = 14    )			->			str:
								UpperCAmelCase     :		Optional[Any]							     =		int(A		,		base=16    )
								UpperCAmelCase     :		Tuple							     =		int(A		,		base=16    )
								UpperCAmelCase     :		Optional[Any]							     =		primes[group]["""prime"""]
								if not DiffieHellman.is_valid_public_key_static(A		,		A    ):
															raise ValueError("""Invalid public key"""    )
								UpperCAmelCase     :		List[str]							     =		pow(A		,		A		,		A    )
								return shaaaa(str(A    ).encode()    ).hexdigest()


if __name__ == "__main__":
							import doctest

							doctest.testmod()






'''simple docstring'''




import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    GPTNeoXForCausalLM,
							    GPTNeoXForQuestionAnswering,
							    GPTNeoXForSequenceClassification,
							    GPTNeoXForTokenClassification,
							    GPTNeoXModel,
							)

class 			UpperCamelCase_ :





	def __init__(  self		,		A		,		A=13		,		A=7		,		A=True		,		A=True		,		A=True		,		A=True		,		A=99		,		A=64		,		A=5		,		A=4		,		A=37		,		A="gelu"		,		A=0.1		,		A=0.1		,		A=512		,		A=16		,		A=2		,		A=0.0_2		,		A=3		,		A=4		,		A=None		,		)			->			Optional[int]:
								UpperCAmelCase     :		List[Any]							     =		parent
								UpperCAmelCase     :		Optional[int]							     =		batch_size
								UpperCAmelCase     :		Union[str, Any]							     =		seq_length
								UpperCAmelCase     :		Optional[Any]							     =		is_training
								UpperCAmelCase     :		Dict							     =		use_input_mask
								UpperCAmelCase     :		str							     =		use_token_type_ids
								UpperCAmelCase     :		List[Any]							     =		use_labels
								UpperCAmelCase     :		List[Any]							     =		vocab_size
								UpperCAmelCase     :		Dict							     =		hidden_size
								UpperCAmelCase     :		Dict							     =		num_hidden_layers
								UpperCAmelCase     :		Optional[int]							     =		num_attention_heads
								UpperCAmelCase     :		int							     =		intermediate_size
								UpperCAmelCase     :		List[str]							     =		hidden_act
								UpperCAmelCase     :		List[str]							     =		hidden_dropout_prob
								UpperCAmelCase     :		int							     =		attention_probs_dropout_prob
								UpperCAmelCase     :		str							     =		max_position_embeddings
								UpperCAmelCase     :		Optional[Any]							     =		type_vocab_size
								UpperCAmelCase     :		List[str]							     =		type_sequence_label_size
								UpperCAmelCase     :		int							     =		initializer_range
								UpperCAmelCase     :		str							     =		num_labels
								UpperCAmelCase     :		Optional[int]							     =		num_choices
								UpperCAmelCase     :		Dict							     =		scope
								UpperCAmelCase     :		Union[str, Any]							     =		vocab_size - 1





	def  _lowercase(  self    )			->			Union[str, Any]:
								UpperCAmelCase     :		List[str]							     =		ids_tensor([self.batch_size, self.seq_length]		,		self.vocab_size    )

								UpperCAmelCase     :		Any							     =		None
								if self.use_input_mask:
															UpperCAmelCase     :		int							     =		random_attention_mask([self.batch_size, self.seq_length]    )

								UpperCAmelCase     :		List[str]							     =		None
								if self.use_labels:
															UpperCAmelCase     :		Optional[int]							     =		ids_tensor([self.batch_size, self.seq_length]		,		self.num_labels    )

								UpperCAmelCase     :		Optional[int]							     =		self.get_config()

								return config, input_ids, input_mask, token_labels





	def  _lowercase(  self    )			->			Optional[Any]:
								return GPTNeoXConfig(
								    vocab_size=self.vocab_size		,		hidden_size=self.hidden_size		,		num_hidden_layers=self.num_hidden_layers		,		num_attention_heads=self.num_attention_heads		,		intermediate_size=self.intermediate_size		,		hidden_act=self.hidden_act		,		hidden_dropout_prob=self.hidden_dropout_prob		,		attention_probs_dropout_prob=self.attention_probs_dropout_prob		,		max_position_embeddings=self.max_position_embeddings		,		type_vocab_size=self.type_vocab_size		,		is_decoder=A		,		initializer_range=self.initializer_range		,		pad_token_id=self.pad_token_id		,		)





	def  _lowercase(  self    )			->			Optional[Any]:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		self.prepare_config_and_inputs()

								UpperCAmelCase     :		Any							     =		True

								return config, input_ids, input_mask, token_labels





	def  _lowercase(  self		,		A		,		A		,		A    )			->			int:
								UpperCAmelCase     :		str							     =		GPTNeoXModel(config=A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		List[str]							     =		model(A		,		attention_mask=A    )
								UpperCAmelCase     :		List[str]							     =		model(A    )
								self.parent.assertEqual(result.last_hidden_state.shape		,		(self.batch_size, self.seq_length, self.hidden_size)    )





	def  _lowercase(  self		,		A		,		A		,		A    )			->			Optional[int]:
								UpperCAmelCase     :		str							     =		True
								UpperCAmelCase     :		Optional[Any]							     =		GPTNeoXModel(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		List[Any]							     =		model(A		,		attention_mask=A    )
								self.parent.assertEqual(result.last_hidden_state.shape		,		(self.batch_size, self.seq_length, self.hidden_size)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			List[str]:
								UpperCAmelCase     :		Tuple							     =		GPTNeoXForCausalLM(config=A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		str							     =		model(A		,		attention_mask=A		,		labels=A    )
								self.parent.assertEqual(result.logits.shape		,		(self.batch_size, self.seq_length, self.vocab_size)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			Tuple:
								UpperCAmelCase     :		List[str]							     =		self.num_labels
								UpperCAmelCase     :		Any							     =		GPTNeoXForQuestionAnswering(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		str							     =		model(A		,		attention_mask=A    )
								self.parent.assertEqual(result.start_logits.shape		,		(self.batch_size, self.seq_length)    )
								self.parent.assertEqual(result.end_logits.shape		,		(self.batch_size, self.seq_length)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			int:
								UpperCAmelCase     :		Tuple							     =		self.num_labels
								UpperCAmelCase     :		List[str]							     =		GPTNeoXForSequenceClassification(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		List[Any]							     =		ids_tensor([self.batch_size]		,		self.type_sequence_label_size    )
								UpperCAmelCase     :		Tuple							     =		model(A		,		attention_mask=A		,		labels=A    )
								self.parent.assertEqual(result.logits.shape		,		(self.batch_size, self.num_labels)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			str:
								UpperCAmelCase     :		List[Any]							     =		self.num_labels
								UpperCAmelCase     :		Tuple							     =		GPTNeoXForTokenClassification(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		int							     =		model(A		,		attention_mask=A		,		labels=A    )
								self.parent.assertEqual(result.logits.shape		,		(self.batch_size, self.seq_length, self.num_labels)    )





	def  _lowercase(  self		,		A		,		A		,		A    )			->			Union[str, Any]:
								UpperCAmelCase     :		Optional[int]							     =		True
								UpperCAmelCase     :		str							     =		GPTNeoXForCausalLM(config=A    )
								model.to(A    )
								model.eval()

								# first forward pass
								UpperCAmelCase     :		List[str]							     =		model(A		,		attention_mask=A		,		use_cache=A    )
								UpperCAmelCase     :		List[Any]							     =		outputs.past_key_values

								# create hypothetical multiple next token and extent to next_input_ids
								UpperCAmelCase     :		Optional[int]							     =		ids_tensor((self.batch_size, 3)		,		config.vocab_size    )
								UpperCAmelCase     :		Any							     =		ids_tensor((self.batch_size, 3)		,		vocab_size=2    )

								# append to next input_ids and
								UpperCAmelCase     :		str							     =		torch.cat([input_ids, next_tokens]		,		dim=-1    )
								UpperCAmelCase     :		Any							     =		torch.cat([input_mask, next_mask]		,		dim=-1    )

								UpperCAmelCase     :		Dict							     =		model(A		,		attention_mask=A		,		output_hidden_states=A    )
								UpperCAmelCase     :		Any							     =		output_from_no_past["""hidden_states"""][0]
								UpperCAmelCase     :		List[str]							     =		model(
								    A		,		attention_mask=A		,		past_key_values=A		,		output_hidden_states=A		,		)["""hidden_states"""][0]

								# select random slice
								UpperCAmelCase     :		Tuple							     =		ids_tensor((1,)		,		output_from_past.shape[-1]    ).item()
								UpperCAmelCase     :		List[Any]							     =		output_from_no_past[:, -3:, random_slice_idx].detach()
								UpperCAmelCase     :		List[str]							     =		output_from_past[:, :, random_slice_idx].detach()

								self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]    )

								# test that outputs are equal for slice
								self.parent.assertTrue(torch.allclose(A		,		A		,		atol=1e-3    )    )






	def  _lowercase(  self    )			->			int:
								UpperCAmelCase     :		Tuple							     =		self.prepare_config_and_inputs()
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Optional[Any]							     =		config_and_inputs
								UpperCAmelCase     :		Union[str, Any]							     =		{"""input_ids""": input_ids, """attention_mask""": input_mask}
								return config, inputs_dict




@require_torch
class 			UpperCamelCase_ (							__magic_name__      ,     __magic_name__      ,     __magic_name__      ,     unittest.TestCase			):
	lowercase          = (
	    (
	        GPTNeoXModel,
	        GPTNeoXForCausalLM,
	        GPTNeoXForQuestionAnswering,
	        GPTNeoXForSequenceClassification,
	        GPTNeoXForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	lowercase          = (GPTNeoXForCausalLM,) if is_torch_available() else ()
	lowercase          = (
	    {
	        'feature-extraction': GPTNeoXModel,
	        'question-answering': GPTNeoXForQuestionAnswering,
	        'text-classification': GPTNeoXForSequenceClassification,
	        'text-generation': GPTNeoXForCausalLM,
	        'token-classification': GPTNeoXForTokenClassification,
	        'zero-shot': GPTNeoXForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowercase          = False
	lowercase          = False
	lowercase          = False
	lowercase          = False





	def  _lowercase(  self    )			->			Union[str, Any]:
								UpperCAmelCase     :		str							     =		GPTNeoXModelTester(self    )
								UpperCAmelCase     :		Optional[Any]							     =		ConfigTester(self		,		config_class=A		,		hidden_size=64		,		num_attention_heads=8    )





	def  _lowercase(  self    )			->			Optional[Any]:
								self.config_tester.run_common_tests()





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Any							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(A		,		A		,		A    )





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Union[str, Any]							     =		self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(A		,		A		,		A    )





	def  _lowercase(  self    )			->			Optional[Any]:
								# This regression test was failing with PyTorch < 1.3
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Optional[Any]							     =		self.model_tester.prepare_config_and_inputs_for_decoder()

								UpperCAmelCase     :		Optional[Any]							     =		None

								self.model_tester.create_and_check_model_as_decoder(A		,		A		,		A    )





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_decoder_model_past_large_inputs(A		,		A		,		A    )





	def  _lowercase(  self    )			->			int:
								UpperCAmelCase     :		Any							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_causal_lm(*A    )





	def  _lowercase(  self    )			->			Optional[int]:
								UpperCAmelCase     :		Tuple							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*A    )





	def  _lowercase(  self    )			->			Any:
								UpperCAmelCase     :		Tuple							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*A    )





	def  _lowercase(  self    )			->			Optional[Any]:
								UpperCAmelCase     :		Optional[Any]							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*A    )





	@unittest.skip(reason="""Feed forward chunking is not implemented"""    )
	def  _lowercase(  self    )			->			Optional[int]:
								pass






	@parameterized.expand([("""linear""",), ("""dynamic""",)]    )
	def  _lowercase(  self		,		A    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     :		List[Any]							     =		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase     :		int							     =		ids_tensor([1, 10]		,		config.vocab_size    )
								UpperCAmelCase     :		Optional[Any]							     =		ids_tensor([1, int(config.max_position_embeddings * 1.5    )]		,		config.vocab_size    )

								set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
								UpperCAmelCase     :		Dict							     =		GPTNeoXModel(A    )
								original_model.to(A    )
								original_model.eval()
								UpperCAmelCase     :		List[str]							     =		original_model(A    ).last_hidden_state
								UpperCAmelCase     :		Any							     =		original_model(A    ).last_hidden_state

								set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
								UpperCAmelCase     :		Any							     =		{"""type""": scaling_type, """factor""": 1_0.0}
								UpperCAmelCase     :		str							     =		GPTNeoXModel(A    )
								scaled_model.to(A    )
								scaled_model.eval()
								UpperCAmelCase     :		Optional[Any]							     =		scaled_model(A    ).last_hidden_state
								UpperCAmelCase     :		Optional[Any]							     =		scaled_model(A    ).last_hidden_state

								# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
								# maximum sequence length, so the outputs for the short input should match.
								if scaling_type == "dynamic":
															self.assertTrue(torch.allclose(A		,		A		,		atol=1e-5    )    )
								else:
															self.assertFalse(torch.allclose(A		,		A		,		atol=1e-5    )    )

								# The output should be different for long inputs
								self.assertFalse(torch.allclose(A		,		A		,		atol=1e-5    )    )




@require_torch
class 			UpperCamelCase_ (							unittest.TestCase			):





	@slow
	def  _lowercase(  self    )			->			List[Any]:
								UpperCAmelCase     :		str							     =		AutoTokenizer.from_pretrained("""EleutherAI/pythia-410m-deduped"""    )
								for checkpointing in [True, False]:
															UpperCAmelCase     :		int							     =		GPTNeoXForCausalLM.from_pretrained("""EleutherAI/pythia-410m-deduped"""    )

															if checkpointing:
																						model.gradient_checkpointing_enable()
															else:
																						model.gradient_checkpointing_disable()
															model.to(A    )

															UpperCAmelCase     :		List[Any]							     =		tokenizer("""My favorite food is"""		,		return_tensors="""pt"""    ).to(A    )
															# The hub repo. is updated on 2023-04-04, resulting in poor outputs.
															# See: https://github.com/huggingface/transformers/pull/24193
															UpperCAmelCase     :		List[str]							     =		"""My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"""

															UpperCAmelCase     :		Union[str, Any]							     =		model.generate(**A		,		do_sample=A		,		max_new_tokens=20    )
															UpperCAmelCase     :		Tuple							     =		tokenizer.batch_decode(A    )[0]

															self.assertEqual(A		,		A    )

