


"""simple docstring"""






from __future__ import annotations

import unittest
import warnings

from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel


@require_tf
class      __lowerCamelCase	:

     '''simple docstring'''
     A_							:						Dict     =		MarianConfig
     A_							:						List[Any]     =		{}
     A_							:						str     =		'gelu'


     def __init__(   self     ,					__UpperCAmelCase     ,					__UpperCAmelCase=13     ,					__UpperCAmelCase=7     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=99     ,					__UpperCAmelCase=32     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=4     ,					__UpperCAmelCase=37     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=20     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=1     ,					__UpperCAmelCase=0     ,					)							->				Optional[Any]:
          _a     						=					parent
          _a     						=					batch_size
          _a     						=					seq_length
          _a     						=					is_training
          _a     						=					use_labels
          _a     						=					vocab_size
          _a     						=					hidden_size
          _a     						=					num_hidden_layers
          _a     						=					num_attention_heads
          _a     						=					intermediate_size

          _a     						=					hidden_dropout_prob
          _a     						=					attention_probs_dropout_prob
          _a     						=					max_position_embeddings
          _a     						=					eos_token_id
          _a     						=					pad_token_id
          _a     						=					bos_token_id


     def 				_UpperCAmelCase   (   self				)							->				Optional[int]:
          _a     						=					ids_tensor([self.batch_size, self.seq_length - 1]     ,					self.vocab_size				)
          _a     						=					tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size				)     ,					1				)
          _a     						=					tf.concat([input_ids, eos_tensor]     ,					axis=1				)

          _a     						=					ids_tensor([self.batch_size, self.seq_length]     ,					self.vocab_size				)

          _a     						=					self.config_cls(
              vocab_size=self.vocab_size     ,					d_model=self.hidden_size     ,					encoder_layers=self.num_hidden_layers     ,					decoder_layers=self.num_hidden_layers     ,					encoder_attention_heads=self.num_attention_heads     ,					decoder_attention_heads=self.num_attention_heads     ,					encoder_ffn_dim=self.intermediate_size     ,					decoder_ffn_dim=self.intermediate_size     ,					dropout=self.hidden_dropout_prob     ,					attention_dropout=self.attention_probs_dropout_prob     ,					max_position_embeddings=self.max_position_embeddings     ,					eos_token_ids=[2]     ,					bos_token_id=self.bos_token_id     ,					pad_token_id=self.pad_token_id     ,					decoder_start_token_id=self.pad_token_id     ,					**self.config_updates     ,					)
          _a     						=					prepare_marian_inputs_dict(__lowerCAmelCase     ,					__lowerCAmelCase     ,					__lowerCAmelCase				)
          return config, inputs_dict





     def 				_UpperCAmelCase   (   self     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)							->				Any:
          _a     						=					TFMarianModel(config=__lowerCAmelCase				).get_decoder()
          _a     						=					inputs_dict['''input_ids''']

          _a     						=					input_ids[:1, :]
          _a     						=					inputs_dict['''attention_mask'''][:1, :]
          _a     						=					inputs_dict['''head_mask''']
          _a     						=					1

          # first forward pass
          _a     						=					model(__lowerCAmelCase     ,					attention_mask=__lowerCAmelCase     ,					head_mask=__lowerCAmelCase     ,					use_cache=__lowerCAmelCase				)

          _a						,   _a     						=					outputs.to_tuple()

          # create hypothetical next token and extent to next_input_ids
          _a     						=					ids_tensor((self.batch_size, 3)     ,					config.vocab_size				)
          _a     						=					tf.cast(ids_tensor((self.batch_size, 3)     ,					2				)     ,					tf.inta				)

          # append to next input_ids and
          _a     						=					tf.concat([input_ids, next_tokens]     ,					axis=-1				)
          _a     						=					tf.concat([attention_mask, next_attn_mask]     ,					axis=-1				)

          _a     						=					model(__lowerCAmelCase     ,					attention_mask=__lowerCAmelCase				)[0]
          _a     						=					model(__lowerCAmelCase     ,					attention_mask=__lowerCAmelCase     ,					past_key_values=__lowerCAmelCase				)[0]

          self.parent.assertEqual(next_tokens.shape[1]     ,					output_from_past.shape[1]				)

          # select random slice
          _a     						=					int(ids_tensor((1,)     ,					output_from_past.shape[-1]				)				)
          _a     						=					output_from_no_past[:, -3:, random_slice_idx]
          _a     						=					output_from_past[:, :, random_slice_idx]

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







def    A_ ( _lowerCAmelCase  :       Union[str, Any],  _lowerCAmelCase  :       Optional[int],  _lowerCAmelCase  :       int,  _lowerCAmelCase  :       Tuple=None,  _lowerCAmelCase  :       Union[str, Any]=None,  _lowerCAmelCase  :       Tuple=None,  _lowerCAmelCase  :       int=None,  _lowerCAmelCase  :       Dict=None,  ):






     """simple docstring"""

     if attention_mask is None:
          _a     						=					tf.cast(tf.math.not_equal(_lowercase,  config.pad_token_id							),  tf.inta							)
     if decoder_attention_mask is None:
          _a     						=					tf.concat(
              [
                  tf.ones(decoder_input_ids[:, :1].shape,  dtype=tf.inta							),
                  tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:],  config.pad_token_id							),  tf.inta							),
              ],  axis=-1,  )
     if head_mask is None:
          _a     						=					tf.ones((config.encoder_layers, config.encoder_attention_heads)							)
     if decoder_head_mask is None:
          _a     						=					tf.ones((config.decoder_layers, config.decoder_attention_heads)							)
     if cross_attn_head_mask is None:
          _a     						=					tf.ones((config.decoder_layers, config.decoder_attention_heads)							)
     return {
         "input_ids": input_ids,
         "decoder_input_ids": decoder_input_ids,
         "attention_mask": attention_mask,
         "decoder_attention_mask": decoder_attention_mask,
         "head_mask": head_mask,
         "decoder_head_mask": decoder_head_mask,
         "cross_attn_head_mask": cross_attn_head_mask,
     }


@require_tf
class      __lowerCamelCase	(       _lowercase	,       _lowercase	,       unittest.TestCase					):

     '''simple docstring'''
     A_							:						List[str]     =		(TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
     A_							:						int     =		(TFMarianMTModel,) if is_tf_available() else ()
     A_							:						Dict     =		(
         {
             'conversational': TFMarianMTModel,
             'feature-extraction': TFMarianModel,
             'summarization': TFMarianMTModel,
             'text2text-generation': TFMarianMTModel,
             'translation': TFMarianMTModel,
         }
         if is_tf_available()
         else {}
     )
     A_							:						List[Any]     =		True
     A_							:						Optional[int]     =		False
     A_							:						Optional[Any]     =		False


     def 				_UpperCAmelCase   (   self				)							->				int:
          _a     						=					TFMarianModelTester(self				)
          _a     						=					ConfigTester(self     ,					config_class=__lowerCAmelCase				)


     def 				_UpperCAmelCase   (   self				)							->				Union[str, Any]:
          self.config_tester.run_common_tests()





     def 				_UpperCAmelCase   (   self				)							->				List[Any]:
          _a     						=					self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.check_decoder_model_past_large_inputs(*__lowerCAmelCase				)







@require_tf
class      __lowerCamelCase	(       unittest.TestCase					):

     '''simple docstring'''
     A_							:						Dict     =		1_000  # show more chars for failing integration tests


     @classmethod
     def 				_UpperCAmelCase   (   cls				)							->				None:
          _a     						=					F'Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}'
          return cls


     @cached_property
     def 				_UpperCAmelCase   (   self				)							->				MarianTokenizer:
          return AutoTokenizer.from_pretrained(self.model_name				)


     @property
     def 				_UpperCAmelCase   (   self				)							->				int:
          return self.tokenizer.eos_token_id


     @cached_property
     def 				_UpperCAmelCase   (   self				)							->				List[Any]:
          warnings.simplefilter('''error'''				)
          _a     						=					TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name				)
          assert isinstance(__lowerCAmelCase     ,					__lowerCAmelCase				)
          _a     						=					model.config
          self.assertListEqual(c.bad_words_ids     ,					[[c.pad_token_id]]				)
          self.assertEqual(c.max_length     ,					512				)
          self.assertEqual(c.decoder_start_token_id     ,					c.pad_token_id				)
          return model


     def 				_UpperCAmelCase   (   self     ,					**__UpperCAmelCase				)							->				int:
          _a     						=					self.translate_src_text(**__lowerCAmelCase				)
          self.assertListEqual(self.expected_text     ,					__lowerCAmelCase				)





     def 				_UpperCAmelCase   (   self     ,					**__UpperCAmelCase				)							->				Union[str, Any]:
          _a     						=					self.tokenizer(self.src_text     ,					**__lowerCAmelCase     ,					padding=__lowerCAmelCase     ,					return_tensors='''tf'''				)
          _a     						=					self.model.generate(
              model_inputs.input_ids     ,					attention_mask=model_inputs.attention_mask     ,					num_beams=2     ,					max_length=128				)
          _a     						=					self.tokenizer.batch_decode(generated_ids.numpy()     ,					skip_special_tokens=__lowerCAmelCase				)
          return generated_words







@require_sentencepiece
@require_tokenizers
@require_tf
class      __lowerCamelCase	(       _lowercase					):

     '''simple docstring'''
     A_							:						List[str]     =		'mt'
     A_							:						Tuple     =		'en'
     A_							:						Tuple     =		['Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.']
     A_							:						Optional[int]     =		['Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.']


     @unittest.skip('''Skipping until #12647 is resolved.'''				)
     @slow
     def 				_UpperCAmelCase   (   self				)							->				Any:
          self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class      __lowerCamelCase	(       _lowercase					):

     '''simple docstring'''
     A_							:						Optional[Any]     =		'en'
     A_							:						Dict     =		'zh'
     A_							:						Tuple     =		['My name is Wolfgang and I live in Berlin']
     A_							:						Optional[Any]     =		['我叫沃尔夫冈 我住在柏林']


     @unittest.skip('''Skipping until #12647 is resolved.'''				)
     @slow
     def 				_UpperCAmelCase   (   self				)							->				Union[str, Any]:
          self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class      __lowerCamelCase	(       _lowercase					):

     '''simple docstring'''
     A_							:						Optional[int]     =		'en'
     A_							:						List[str]     =		'ROMANCE'
     A_							:						int     =		[
         '>>fr<< Don\'t spend so much time watching TV.',
         '>>pt<< Your message has been sent.',
         '>>es<< He\'s two years older than me.',
     ]
     A_							:						str     =		[
         'Ne passez pas autant de temps à regarder la télé.',
         'A sua mensagem foi enviada.',
         'Es dos años más viejo que yo.',
     ]


     @unittest.skip('''Skipping until #12647 is resolved.'''				)
     @slow
     def 				_UpperCAmelCase   (   self				)							->				int:
          self._assert_generated_batch_equal_expected()





     @unittest.skip('''Skipping until #12647 is resolved.'''				)
     @slow
     def 				_UpperCAmelCase   (   self				)							->				int:
          _a     						=					TranslationPipeline(self.model     ,					self.tokenizer     ,					framework='''tf'''				)
          _a     						=					pipeline(self.src_text				)
          self.assertEqual(self.expected_text     ,					[x['''translation_text'''] for x in output]				)






import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin




class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = UNetaDModel
					__A        = "sample"




					@property
					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						4
							UpperCamelCase__								=						14
							UpperCamelCase__								=						16

							UpperCamelCase__								=						floats_tensor((batch_size, num_features, seq_len)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([10] * batch_size							).to(__lowerCAmelCase							)

							return {"sample": noise, "timestep": time_step}




					@property
					def 						snake_case_		(    self							)    ->	Optional[int]:
							return (4, 14, 16)




					@property
					def 						snake_case_		(    self							)    ->	Any:
							return (4, 14, 16)




					def 						snake_case_		(    self							)    ->	Any:
							pass




					def 						snake_case_		(    self							)    ->	Tuple:
							pass




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Any:
							super().test_determinism()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	int:
							super().test_outputs_equivalence()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_from_save_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Dict:
							super().test_output()




					def 						snake_case_		(    self							)    ->	int:
							UpperCamelCase__								=						{
							    """block_out_channels""": (32, 64, 128, 256),
							    """in_channels""": 14,
							    """out_channels""": 14,
							    """time_embedding_type""": """positional""",
							    """use_timestep_embedding""": True,
							    """flip_sin_to_cos""": False,
							    """freq_shift""": 1.0,
							    """out_block_type""": """OutConv1DBlock""",
							    """mid_block_type""": """MidResTemporalBlock1D""",
							    """down_block_types""": ("""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""),
							    """up_block_types""": ("""UpResnetBlock1D""", """UpResnetBlock1D""", """UpResnetBlock1D"""),
							    """act_fn""": """mish""",
							}
							UpperCamelCase__								=						self.dummy_input
							return init_dict, inputs_dict




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__	,      UpperCamelCase__								=						UNetaDModel.from_pretrained(
							    """bglick13/hopper-medium-v2-value-function-hor32"""     ,   output_loading_info=__lowerCAmelCase     ,   subfolder="""unet"""							)
							self.assertIsNotNone(__lowerCAmelCase							)
							self.assertEqual(len(loading_info["""missing_keys"""]							)     ,   0							)

							model.to(__lowerCAmelCase							)
							UpperCamelCase__								=						model(**self.dummy_input							)

							assert image is not None, "Make sure output is not None"




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						UNetaDModel.from_pretrained("""bglick13/hopper-medium-v2-value-function-hor32"""     ,   subfolder="""unet"""							)
							torch.manual_seed(0							)
							if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

							UpperCamelCase__								=						model.config.in_channels
							UpperCamelCase__								=						16
							UpperCamelCase__								=						torch.randn((1, seq_len, num_features)							).permute(
							    0     ,   2     ,   1							)  # match original, we can update values and remove
							UpperCamelCase__								=						torch.full((num_features,)     ,   0							)

							with torch.no_grad():
									UpperCamelCase__								=						model(__lowerCAmelCase     ,   __lowerCAmelCase							).sample.permute(0     ,   2     ,   1							)

							UpperCamelCase__								=						output[0, -3:, -3:].flatten()
							# fmt: off
							UpperCamelCase__								=						torch.tensor([-2.1_3_7_1_7_2, 1.1_4_2_6_0_1_6, 0.3_6_8_8_6_8_7, -0.7_6_6_9_2_2, 0.7_3_0_3_1_4_6, 0.1_1_0_3_8_8_6_4, -0.4_7_6_0_6_3_3, 0.1_3_2_7_0_1_7_2, 0.0_2_5_9_1_3_4_8]							)
							# fmt: on
							self.assertTrue(torch.allclose(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-3							)							)




					def 						snake_case_		(    self							)    ->	int:
							# Not implemented yet for this UNet
							pass




					@slow
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""harmonai/maestro-150k"""
							UpperCamelCase__								=						UNetaDModel.from_pretrained(__lowerCAmelCase     ,   subfolder="""unet"""							)
							model.to(__lowerCAmelCase							)

							UpperCamelCase__								=						65_536
							UpperCamelCase__								=						torch.sin(torch.arange(__lowerCAmelCase							)[None, None, :].repeat(1     ,   2     ,   1							)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([1]							).to(__lowerCAmelCase							)

							with torch.no_grad():
									UpperCamelCase__								=						model(__lowerCAmelCase     ,   __lowerCAmelCase							).sample

							UpperCamelCase__								=						output.abs().sum()
							UpperCamelCase__								=						output.abs().max()

							assert (output_sum - 2_2_4.0_8_9_6).abs() < 0.5
							assert (output_max - 0.0_6_0_7).abs() < 4E-4




class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = UNetaDModel
					__A        = "sample"




					@property
					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						4
							UpperCamelCase__								=						14
							UpperCamelCase__								=						16

							UpperCamelCase__								=						floats_tensor((batch_size, num_features, seq_len)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([10] * batch_size							).to(__lowerCAmelCase							)

							return {"sample": noise, "timestep": time_step}




					@property
					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							return (4, 14, 16)




					@property
					def 						snake_case_		(    self							)    ->	List[Any]:
							return (4, 14, 1)




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Dict:
							super().test_determinism()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_outputs_equivalence()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_from_save_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[str]:
							# UNetRL is a value-function is different output shape
							UpperCamelCase__	,      UpperCamelCase__								=						self.prepare_init_args_and_inputs_for_common()
							UpperCamelCase__								=						self.model_class(**__lowerCAmelCase							)
							model.to(__lowerCAmelCase							)
							model.eval()

							with torch.no_grad():
									UpperCamelCase__								=						model(**__lowerCAmelCase							)

									if isinstance(__lowerCAmelCase     ,   __lowerCAmelCase							):
											UpperCamelCase__								=						output.sample

							self.assertIsNotNone(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.Size((inputs_dict["""sample"""].shape[0], 1)							)
							self.assertEqual(output.shape     ,   __lowerCAmelCase     ,   """Input and output shapes do not match"""							)




					def 						snake_case_		(    self							)    ->	Optional[int]:
							pass




					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							pass




					def 						snake_case_		(    self							)    ->	Dict:
							UpperCamelCase__								=						{
							    """in_channels""": 14,
							    """out_channels""": 14,
							    """down_block_types""": ["""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""],
							    """up_block_types""": [],
							    """out_block_type""": """ValueFunction""",
							    """mid_block_type""": """ValueFunctionMidBlock1D""",
							    """block_out_channels""": [32, 64, 128, 256],
							    """layers_per_block""": 1,
							    """downsample_each_block""": True,
							    """use_timestep_embedding""": True,
							    """freq_shift""": 1.0,
							    """flip_sin_to_cos""": False,
							    """time_embedding_type""": """positional""",
							    """act_fn""": """mish""",
							}
							UpperCamelCase__								=						self.dummy_input
							return init_dict, inputs_dict




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__	,      UpperCamelCase__								=						UNetaDModel.from_pretrained(
							    """bglick13/hopper-medium-v2-value-function-hor32"""     ,   output_loading_info=__lowerCAmelCase     ,   subfolder="""value_function"""							)
							self.assertIsNotNone(__lowerCAmelCase							)
							self.assertEqual(len(vf_loading_info["""missing_keys"""]							)     ,   0							)

							value_function.to(__lowerCAmelCase							)
							UpperCamelCase__								=						value_function(**self.dummy_input							)

							assert image is not None, "Make sure output is not None"




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__	,      UpperCamelCase__								=						UNetaDModel.from_pretrained(
							    """bglick13/hopper-medium-v2-value-function-hor32"""     ,   output_loading_info=__lowerCAmelCase     ,   subfolder="""value_function"""							)
							torch.manual_seed(0							)
							if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

							UpperCamelCase__								=						value_function.config.in_channels
							UpperCamelCase__								=						14
							UpperCamelCase__								=						torch.randn((1, seq_len, num_features)							).permute(
							    0     ,   2     ,   1							)  # match original, we can update values and remove
							UpperCamelCase__								=						torch.full((num_features,)     ,   0							)

							with torch.no_grad():
									UpperCamelCase__								=						value_function(__lowerCAmelCase     ,   __lowerCAmelCase							).sample

							# fmt: off
							UpperCamelCase__								=						torch.tensor([1_6_5.2_5] * seq_len							)
							# fmt: on
							self.assertTrue(torch.allclose(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-3							)							)




					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							# Not implemented yet for this UNet
							pass





