
"""simple docstring"""


from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers.models.gptj.modeling_tf_gptj import (
						    TFGPTJForCausalLM,
						    TFGPTJForQuestionAnswering,
						    TFGPTJForSequenceClassification,
						    TFGPTJModel,
						    shape_list,
						)




class    __lowerCamelCase     :




				def __init__(							self       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     parent
							lowerCamelCase_ 		=     13
							lowerCamelCase_ 		=     7
							lowerCamelCase_ 		=     True
							lowerCamelCase_ 		=     True
							lowerCamelCase_ 		=     True
							lowerCamelCase_ 		=     True
							lowerCamelCase_ 		=     True
							lowerCamelCase_ 		=     99
							lowerCamelCase_ 		=     32
							lowerCamelCase_ 		=     4
							lowerCamelCase_ 		=     2
							lowerCamelCase_ 		=     4
							lowerCamelCase_ 		=     37
							lowerCamelCase_ 		=     '''gelu'''
							lowerCamelCase_ 		=     0.1
							lowerCamelCase_ 		=     0.1
							lowerCamelCase_ 		=     512
							lowerCamelCase_ 		=     16
							lowerCamelCase_ 		=     2
							lowerCamelCase_ 		=     0.0_2
							lowerCamelCase_ 		=     3
							lowerCamelCase_ 		=     4
							lowerCamelCase_ 		=     None
							lowerCamelCase_ 		=     self.vocab_size - 1
							lowerCamelCase_ 		=     self.vocab_size - 1
							lowerCamelCase_ 		=     self.vocab_size - 1




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size      )

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

							lowerCamelCase_ 		=     None
							if self.use_token_type_ids:
										lowerCamelCase_ 		=     ids_tensor([self.batch_size, self.seq_length]       ,	self.type_vocab_size      )

							lowerCamelCase_ 		=     None
							if self.use_mc_token_ids:
										lowerCamelCase_ 		=     ids_tensor([self.batch_size, self.num_choices]       ,	self.seq_length      )

							lowerCamelCase_ 		=     None
							lowerCamelCase_ 		=     None
							lowerCamelCase_ 		=     None
							if self.use_labels:
										lowerCamelCase_ 		=     ids_tensor([self.batch_size]       ,	self.type_sequence_label_size      )
										lowerCamelCase_ 		=     ids_tensor([self.batch_size, self.seq_length]       ,	self.num_labels      )
										lowerCamelCase_ 		=     ids_tensor([self.batch_size]       ,	self.num_choices      )

							lowerCamelCase_ 		=     GPTJConfig(
							    vocab_size=self.vocab_size       ,	n_embd=self.hidden_size       ,	n_layer=self.num_hidden_layers       ,	n_head=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       ,	n_positions=self.max_position_embeddings       ,	type_vocab_size=self.type_vocab_size       ,	initializer_range=self.initializer_range       ,	bos_token_id=self.bos_token_id       ,	eos_token_id=self.eos_token_id       ,	pad_token_id=self.pad_token_id       ,	rotary_dim=self.rotary_dim       ,	return_dict=UpperCAmelCase       ,	)

							lowerCamelCase_ 		=     ids_tensor([self.num_hidden_layers, self.num_attention_heads]       ,	2      )

							return (
							    config,
							    input_ids,
							    input_mask,
							    head_mask,
							    token_type_ids,
							    mc_token_ids,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							)




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	*UpperCAmelCase      ):
							lowerCamelCase_ 		=     TFGPTJModel(config=UpperCAmelCase      )
							lowerCamelCase_ 		=     {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''token_type_ids''': token_type_ids,
							}
							lowerCamelCase_ 		=     model(UpperCAmelCase      )

							lowerCamelCase_ 		=     [input_ids, None, input_mask]  # None is the input for 'past'
							lowerCamelCase_ 		=     model(UpperCAmelCase      )

							lowerCamelCase_ 		=     model(UpperCAmelCase      )

							self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	*UpperCAmelCase      ):
							lowerCamelCase_ 		=     TFGPTJModel(config=UpperCAmelCase      )

							# first forward pass
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	token_type_ids=UpperCAmelCase       ,	use_cache=UpperCAmelCase      )
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	token_type_ids=UpperCAmelCase      )
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	token_type_ids=UpperCAmelCase       ,	use_cache=UpperCAmelCase      )

							self.parent.assertTrue(len(UpperCAmelCase      ) == len(UpperCAmelCase      )      )
							self.parent.assertTrue(len(UpperCAmelCase      ) == len(UpperCAmelCase      ) + 1      )

							lowerCamelCase_		,   lowerCamelCase_ 		=     outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							lowerCamelCase_ 		=     ids_tensor((self.batch_size, 1)       ,	config.vocab_size      )
							lowerCamelCase_ 		=     ids_tensor([self.batch_size, 1]       ,	self.type_vocab_size      )

							# append to next input_ids and token_type_ids
							lowerCamelCase_ 		=     tf.concat([input_ids, next_tokens]       ,	axis=-1      )
							lowerCamelCase_ 		=     tf.concat([token_type_ids, next_token_types]       ,	axis=-1      )

							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	token_type_ids=UpperCAmelCase      )['''last_hidden_state''']
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	token_type_ids=UpperCAmelCase       ,	past_key_values=UpperCAmelCase      )[
							    '''last_hidden_state'''
							]

							# select random slice
							lowerCamelCase_ 		=     int(ids_tensor((1,)       ,	shape_list(UpperCAmelCase      )[-1]      )      )
							lowerCamelCase_ 		=     output_from_no_past[:, -1, random_slice_idx]
							lowerCamelCase_ 		=     output_from_past[:, 0, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(UpperCAmelCase       ,	UpperCAmelCase       ,	rtol=1e-6      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	*UpperCAmelCase      ):
							lowerCamelCase_ 		=     TFGPTJModel(config=UpperCAmelCase      )

							# create attention mask
							lowerCamelCase_ 		=     self.seq_length // 2
							lowerCamelCase_ 		=     tf.ones((self.batch_size, half_seq_length)       ,	dtype=tf.intaa      )
							lowerCamelCase_ 		=     tf.zeros((self.batch_size, self.seq_length - half_seq_length)       ,	dtype=tf.intaa      )
							lowerCamelCase_ 		=     tf.concat([attn_mask_begin, attn_mask_end]       ,	axis=1      )

							# first forward pass
							lowerCamelCase_		,   lowerCamelCase_ 		=     model(UpperCAmelCase       ,	attention_mask=UpperCAmelCase      ).to_tuple()

							# create hypothetical next token and extent to next_input_ids
							lowerCamelCase_ 		=     ids_tensor((self.batch_size, 1)       ,	config.vocab_size      )

							# change a random masked slice from input_ids
							lowerCamelCase_ 		=     ids_tensor((1,)       ,	UpperCAmelCase      ).numpy() + 1
							lowerCamelCase_ 		=     ids_tensor((self.batch_size, self.seq_length)       ,	config.vocab_size      )
							lowerCamelCase_ 		=     tf.range(self.seq_length      ) == (self.seq_length - random_seq_idx_to_change)
							lowerCamelCase_ 		=     tf.transpose(
							    tf.broadcast_to(tf.expand_dims(UpperCAmelCase       ,	-1      )       ,	(self.seq_length, self.batch_size)      )      )
							lowerCamelCase_ 		=     tf.where(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      )

							# append to next input_ids and attn_mask
							lowerCamelCase_ 		=     tf.concat([input_ids, next_tokens]       ,	axis=-1      )
							lowerCamelCase_ 		=     tf.concat([attn_mask, tf.ones((shape_list(UpperCAmelCase      )[0], 1)       ,	dtype=tf.intaa      )]       ,	axis=1      )

							# get two different outputs
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	attention_mask=UpperCAmelCase      )['''last_hidden_state''']
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	past_key_values=UpperCAmelCase       ,	attention_mask=UpperCAmelCase      )[
							    '''last_hidden_state'''
							]

							# select random slice
							lowerCamelCase_ 		=     int(ids_tensor((1,)       ,	shape_list(UpperCAmelCase      )[-1]      )      )
							lowerCamelCase_ 		=     output_from_no_past[:, -1, random_slice_idx]
							lowerCamelCase_ 		=     output_from_past[:, 0, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(UpperCAmelCase       ,	UpperCAmelCase       ,	rtol=1e-1_2      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	*UpperCAmelCase      ):
							lowerCamelCase_ 		=     TFGPTJModel(config=UpperCAmelCase      )

							lowerCamelCase_ 		=     input_ids[:1, :]
							lowerCamelCase_ 		=     input_mask[:1, :]
							lowerCamelCase_ 		=     token_type_ids[:1, :]
							lowerCamelCase_ 		=     1

							# first forward pass
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	attention_mask=UpperCAmelCase       ,	token_type_ids=UpperCAmelCase       ,	use_cache=UpperCAmelCase      )

							lowerCamelCase_		,   lowerCamelCase_ 		=     outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							lowerCamelCase_ 		=     ids_tensor((self.batch_size, 3)       ,	config.vocab_size      )
							lowerCamelCase_ 		=     ids_tensor((self.batch_size, 3)       ,	2      )
							lowerCamelCase_ 		=     ids_tensor((self.batch_size, 3)       ,	self.type_vocab_size      )

							# append to next input_ids and token_type_ids
							lowerCamelCase_ 		=     tf.concat([input_ids, next_tokens]       ,	axis=-1      )
							lowerCamelCase_ 		=     tf.concat([input_mask, next_attn_mask]       ,	axis=-1      )
							lowerCamelCase_ 		=     tf.concat([token_type_ids, next_token_types]       ,	axis=-1      )

							lowerCamelCase_ 		=     model(
							    UpperCAmelCase       ,	token_type_ids=UpperCAmelCase       ,	attention_mask=UpperCAmelCase      )['''last_hidden_state''']
							lowerCamelCase_ 		=     model(
							    UpperCAmelCase       ,	token_type_ids=UpperCAmelCase       ,	attention_mask=UpperCAmelCase       ,	past_key_values=UpperCAmelCase       ,	)['''last_hidden_state''']
							self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]      )

							# select random slice
							lowerCamelCase_ 		=     int(ids_tensor((1,)       ,	shape_list(UpperCAmelCase      )[-1]      )      )
							lowerCamelCase_ 		=     output_from_no_past[:, -3:, random_slice_idx]
							lowerCamelCase_ 		=     output_from_past[:, :, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(UpperCAmelCase       ,	UpperCAmelCase       ,	rtol=1e-3      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	*UpperCAmelCase      ):
							lowerCamelCase_ 		=     TFGPTJForCausalLM(config=UpperCAmelCase      )
							lowerCamelCase_ 		=     {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''token_type_ids''': token_type_ids,
							}
							lowerCamelCase_ 		=     model(UpperCAmelCase      )
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)      )





				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.prepare_config_and_inputs()

							(
							    (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,   (
							    lowerCamelCase_   
							)		,      
							) 		=     config_and_inputs

							lowerCamelCase_ 		=     {
							    '''input_ids''': input_ids,
							    '''token_type_ids''': token_type_ids,
							    '''attention_mask''': input_mask,
							}
							return config, inputs_dict




@require_tf
class    __lowerCamelCase     (   lowerCAmelCase	,  lowerCAmelCase	,  lowerCAmelCase	,  unittest.TestCase ):
				a__:			Dict									=						(
				    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
				    if is_tf_available()
				    else ()
				)

				a__:			str									=						(TFGPTJForCausalLM,) if is_tf_available() else ()
				a__:			str									=						(
				    {
				        'feature-extraction': TFGPTJModel,
				        'question-answering': TFGPTJForQuestionAnswering,
				        'text-classification': TFGPTJForSequenceClassification,
				        'text-generation': TFGPTJForCausalLM,
				        'zero-shot': TFGPTJForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				a__:			Optional[int]									=						False
				a__:			List[str]									=						False
				a__:			str									=						False
				a__:			Optional[Any]									=						False




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							if (
							    pipeline_test_casse_name == "QAPipelineTests"
							    and tokenizer_name is not None
							    and not tokenizer_name.endswith('''Fast'''      )
							):
										# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
										# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
										# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
										return True

							return False




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     TFGPTJModelTester(self      )
							lowerCamelCase_ 		=     ConfigTester(self       ,	config_class=UpperCAmelCase       ,	n_embd=37      )




				def 					UpperCAmelCase__	(							self      ):
							self.config_tester.run_common_tests()




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model(*UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model_past(*UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model_attention_mask_past(*UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model_past_large_inputs(*UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_lm_head_model(*UpperCAmelCase      )




				@slow
				@unittest.skipIf(
				    not is_tf_available() or len(tf.config.list_physical_devices('''GPU'''      )      ) > 0       ,	'''skip testing on GPU for now to avoid GPU OOM.'''       ,	)
				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     TFGPTJModel.from_pretrained('''EleutherAI/gpt-j-6B'''       ,	from_pt=UpperCAmelCase      )
							self.assertIsNotNone(UpperCAmelCase      )





				@unittest.skip(reason='''Currently, model embeddings are going to undergo a major refactor.'''      )
				def 					UpperCAmelCase__	(							self      ):
							super().test_resize_token_embeddings()




@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class    __lowerCamelCase     (   unittest.TestCase ):




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''       ,	from_pt=UpperCAmelCase      )
							lowerCamelCase_ 		=     tf.convert_to_tensor([[464, 3290]]       ,	dtype=tf.intaa      )  # The dog
							# fmt: off
							# The dog is a man's best friend. It is a loyal companion, and it is a friend
							lowerCamelCase_ 		=     [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 1_5185, 11, 290, 340, 318, 257, 1545]
							# fmt: on
							lowerCamelCase_ 		=     model.generate(UpperCAmelCase       ,	do_sample=UpperCAmelCase      )
							self.assertListEqual(output_ids[0].numpy().tolist()       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''       ,	revision='''float16'''      )
							lowerCamelCase_ 		=     TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''       ,	revision='''float16'''       ,	from_pt=UpperCAmelCase      )

							lowerCamelCase_ 		=     tokenizer('''Today is a nice day and'''       ,	return_tensors='''tf'''      )
							# forces the generation to happen on CPU, to avoid GPU-related quirks
							with tf.device(''':/CPU:0'''      ):
										lowerCamelCase_ 		=     model.generate(**UpperCAmelCase       ,	do_sample=UpperCAmelCase       ,	seed=[42, 0]      )
							lowerCamelCase_ 		=     tokenizer.decode(output_ids[0]       ,	skip_special_tokens=UpperCAmelCase      )

							lowerCamelCase_ 		=     '''Today is a nice day and I’m going to go for a walk. I’'''
							self.assertEqual(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''       ,	revision='''float16'''       ,	from_pt=UpperCAmelCase      )
							lowerCamelCase_ 		=     AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''       ,	revision='''float16'''      )

							lowerCamelCase_ 		=     '''left'''

							# Define PAD Token = EOS Token = 50256
							lowerCamelCase_ 		=     tokenizer.eos_token
							lowerCamelCase_ 		=     model.config.eos_token_id

							# use different length sentences to test batching
							lowerCamelCase_ 		=     [
							    '''Hello, my dog is a little''',
							    '''Today, I''',
							]
							lowerCamelCase_ 		=     [
							    '''Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia''',
							    '''Today, I’m going to be talking about a topic that’''',
							]
							return model, tokenizer, sentences, expected_output_sentences




				def 					UpperCAmelCase__	(							self      ):
							# Confirms that we get the expected results with left-padded beam search
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     self._get_beam_search_test_objects()

							lowerCamelCase_ 		=     tokenizer(UpperCAmelCase       ,	return_tensors='''tf'''       ,	padding=UpperCAmelCase      )
							lowerCamelCase_ 		=     model.generate(**UpperCAmelCase       ,	do_sample=UpperCAmelCase       ,	num_beams=2      )
							lowerCamelCase_ 		=     tokenizer.batch_decode(UpperCAmelCase       ,	skip_special_tokens=UpperCAmelCase      )
							self.assertListEqual(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# Confirms that left-padding is working properly
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     self._get_beam_search_test_objects()

							lowerCamelCase_ 		=     tokenizer(UpperCAmelCase       ,	return_tensors='''tf'''       ,	padding=UpperCAmelCase      )
							lowerCamelCase_ 		=     tokenizer(sentences[0]       ,	return_tensors='''tf'''      )
							lowerCamelCase_ 		=     model.generate(**UpperCAmelCase       ,	do_sample=UpperCAmelCase       ,	num_beams=2      )
							lowerCamelCase_ 		=     (
							    shape_list(inputs_non_padded['''input_ids''']      )[-1]
							    - tf.reduce_sum(tf.cast(inputs['''attention_mask'''][-1]       ,	tf.intaa      )      ).numpy()
							)
							lowerCamelCase_ 		=     tokenizer(sentences[1]       ,	return_tensors='''tf'''      )
							lowerCamelCase_ 		=     model.generate(
							    **UpperCAmelCase       ,	do_sample=UpperCAmelCase       ,	num_beams=2       ,	max_length=model.config.max_length - num_paddings      )
							lowerCamelCase_ 		=     tokenizer.decode(output_non_padded[0]       ,	skip_special_tokens=UpperCAmelCase      )
							lowerCamelCase_ 		=     tokenizer.decode(output_padded[0]       ,	skip_special_tokens=UpperCAmelCase      )
							self.assertListEqual(UpperCAmelCase       ,	[non_padded_sentence, padded_sentence]      )





				def 					UpperCAmelCase__	(							self      ):
							# Confirms that XLA is working properly
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     self._get_beam_search_test_objects()

							lowerCamelCase_ 		=     tokenizer(UpperCAmelCase       ,	return_tensors='''tf'''       ,	padding=UpperCAmelCase      )
							lowerCamelCase_ 		=     tf.function(model.generate       ,	jit_compile=UpperCAmelCase      )
							lowerCamelCase_ 		=     xla_generate(**UpperCAmelCase       ,	do_sample=UpperCAmelCase       ,	num_beams=2      )
							lowerCamelCase_ 		=     tokenizer.batch_decode(UpperCAmelCase       ,	skip_special_tokens=UpperCAmelCase      )
							self.assertListEqual(UpperCAmelCase       ,	UpperCAmelCase      )








"""simple docstring"""


import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()




class    __lowerCamelCase     (   unittest.TestCase ):
				a__:			Union[str, Any]									=						'hf-internal-testing/tiny-stable-diffusion-pipe'
				a__:			Tuple									=						1
				a__:			int									=						77
				a__:			List[str]									=						32
				a__:			List[Any]									=						4
				a__:			str									=						64
				a__:			Optional[Any]									=						torch.manual_seed(0 )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase=0.9_9_9_9      ):
							lowerCamelCase_ 		=     UNetaDConditionModel.from_pretrained(self.model_id       ,	subfolder='''unet'''      )
							lowerCamelCase_ 		=     unet.to(UpperCAmelCase      )
							lowerCamelCase_ 		=     EMAModel(unet.parameters()       ,	decay=UpperCAmelCase       ,	model_cls=UpperCAmelCase       ,	model_config=unet.config      )
							return unet, ema_unet




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     torch.randn(
							    self.batch_size       ,	self.num_in_channels       ,	self.latent_height       ,	self.latent_width       ,	generator=self.generator      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.randint(0       ,	1000       ,	size=(self.batch_size,)       ,	generator=self.generator      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.randn(
							    self.batch_size       ,	self.prompt_length       ,	self.text_encoder_hidden_dim       ,	generator=self.generator      ).to(UpperCAmelCase      )
							return noisy_latents, timesteps, encoder_hidden_states




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     {}
							for k, param in unet.state_dict().items():
										lowerCamelCase_ 		=     torch.randn_like(UpperCAmelCase      ) + (param * torch.randn_like(UpperCAmelCase      ))
										updated_state_dict.update({k: updated_param}      )
							unet.load_state_dict(UpperCAmelCase      )
							return unet




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Take the first (hypothetical) EMA step.
							ema_unet.step(unet.parameters()      )
							assert ema_unet.optimization_step == 1

							# Take two more.
							for _ in range(2      ):
										ema_unet.step(unet.parameters()      )
							assert ema_unet.optimization_step == 3




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Since the `unet` is not being updated (i.e., backprop'd)
							# there won't be any difference between the `params` of `unet`
							# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
							ema_unet.step(unet.parameters()      )
							lowerCamelCase_ 		=     list(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )

							# The above holds true even if we call `ema.step()` multiple times since
							# `unet` params are still not being updated.
							for _ in range(4      ):
										ema_unet.step(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Here we simulate the parameter updates for `unet`. Since there might
							# be some parameters which are initialized to zero we take extra care to
							# initialize their values to something non-zero before the multiplication.
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )

							# Take the EMA step.
							ema_unet.step(unet_pseudo_updated_step_one.parameters()      )

							# Now the EMA'd parameters won't be equal to the original model parameters.
							lowerCamelCase_ 		=     list(unet_pseudo_updated_step_one.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )

							# Ensure this is the case when we take multiple EMA steps.
							for _ in range(4      ):
										ema_unet.step(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# If we call EMA step after a backpropagation consecutively for two times,
							# the shadow params from those two steps should be different.
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()

							# First backprop + EMA
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_one.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							# Second backprop + EMA
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_two.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							for step_one, step_two in zip(UpperCAmelCase       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# If there's no decay even if there are backprops, EMA steps
							# won't take any effect i.e., the shadow params would remain the
							# same.
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models(decay=0.0      )
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_one.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_two.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							for step_one, step_two in zip(UpperCAmelCase       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )





				@skip_mps
				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_dummy_inputs()

							with tempfile.TemporaryDirectory() as tmpdir:
										ema_unet.save_pretrained(UpperCAmelCase      )
										lowerCamelCase_ 		=     UNetaDConditionModel.from_pretrained(UpperCAmelCase       ,	model_cls=UpperCAmelCase      )
										lowerCamelCase_ 		=     loaded_unet.to(unet.device      )

							# Since no EMA step has been performed the outputs should match.
							lowerCamelCase_ 		=     unet(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ).sample
							lowerCamelCase_ 		=     loaded_unet(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ).sample

							assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase       ,	atol=1e-4      )



