





import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
		import jax
		import jax.numpy as jnp
		import numpy as np

		from transformers.modeling_flax_pytorch_utils import (
		    convert_pytorch_state_dict_to_flax,
		    load_flax_weights_in_pytorch_model,
		)
		from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
		import torch

@require_flax
class 		snake_case_      :







	'''simple docstring'''




	def __init__(	self       :    Optional[int]				,			__lowerCamelCase       :    Union[str, Any]				,			__lowerCamelCase       :    Tuple=14				,			__lowerCamelCase       :    str=7				,			__lowerCamelCase       :    str=True				,			__lowerCamelCase       :    str=True				,			__lowerCamelCase       :    List[str]=True				,			__lowerCamelCase       :    Dict=99				,			__lowerCamelCase       :    Any=32				,			__lowerCamelCase       :    List[str]=5				,			__lowerCamelCase       :    Union[str, Any]=4				,			__lowerCamelCase       :    Any=37				,			__lowerCamelCase       :    Optional[int]="gelu"				,			__lowerCamelCase       :    str=0.1				,			__lowerCamelCase       :    Union[str, Any]=0.1				,			__lowerCamelCase       :    Any=512				,			__lowerCamelCase       :    Dict=0.02				,			__lowerCamelCase       :    Dict=None				,			)					->   Dict:

			'''simple docstring'''




			__lowercase					=    parent
			__lowercase					=    batch_size
			__lowercase					=    seq_length
			__lowercase					=    is_training
			__lowercase					=    use_input_mask
			__lowercase					=    use_labels
			__lowercase					=    vocab_size
			__lowercase					=    d_model
			__lowercase					=    num_hidden_layers
			__lowercase					=    num_attention_heads
			__lowercase					=    ffn_dim
			__lowercase					=    activation_function
			__lowercase					=    activation_dropout
			__lowercase					=    attention_dropout
			__lowercase					=    max_position_embeddings
			__lowercase					=    initializer_range
			__lowercase					=    None
			__lowercase					=    0
			__lowercase					=    2
			__lowercase					=    1




	def       UpperCAmelCase (	self       :    List[Any]      )					->   Optional[int]:

			'''simple docstring'''




			__lowercase					=    np.clip(ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size      )				,			3				,			self.vocab_size      )

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

			__lowercase					=    XGLMConfig(
			    vocab_size=self.vocab_size				,			d_model=self.hidden_size				,			num_layers=self.num_hidden_layers				,			attention_heads=self.num_attention_heads				,			ffn_dim=self.ffn_dim				,			activation_function=self.activation_function				,			activation_dropout=self.activation_dropout				,			attention_dropout=self.attention_dropout				,			max_position_embeddings=self.max_position_embeddings				,			initializer_range=self.initializer_range				,			use_cache=__lowerCamelCase				,			bos_token_id=self.bos_token_id				,			eos_token_id=self.eos_token_id				,			pad_token_id=self.pad_token_id				,			)

			return (config, input_ids, input_mask)




	def       UpperCAmelCase (	self       :    Optional[Any]      )					->   Union[str, Any]:

			'''simple docstring'''




			__lowercase					=    self.prepare_config_and_inputs()
			__lowercase	,			__lowercase	,			__lowercase					=    config_and_inputs
			__lowercase					=    {'input_ids': input_ids, 'attention_mask': attention_mask}
			return config, inputs_dict




	def       UpperCAmelCase (	self       :    List[Any]      )					->   str:

			'''simple docstring'''




			__lowercase	,			__lowercase	,			__lowercase					=    self.prepare_config_and_inputs()

			__lowercase					=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
			__lowercase					=    ids_tensor([self.batch_size, self.seq_length]				,			vocab_size=2      )

			return (
			    config,
			    input_ids,
			    attention_mask,
			    encoder_hidden_states,
			    encoder_attention_mask,
			)




	def       UpperCAmelCase (	self       :    Optional[int]				,			__lowerCamelCase       :    Optional[int]				,			__lowerCamelCase       :    Tuple				,			__lowerCamelCase       :    Optional[int]				,			__lowerCamelCase       :    List[Any]      )					->   Any:

			'''simple docstring'''




			__lowercase					=    20
			__lowercase					=    model_class_name(__lowerCamelCase      )

			__lowercase					=    model.init_cache(input_ids.shape[0]				,			__lowerCamelCase      )
			__lowercase					=    jnp.ones((input_ids.shape[0], max_decoder_length)				,			dtype='i4'      )

			__lowercase					=    jnp.broadcast_to(
			    jnp.arange(input_ids.shape[-1] - 1      )[None, :]				,			(input_ids.shape[0], input_ids.shape[-1] - 1)      )
			__lowercase					=    model(
			    input_ids[:, :-1]				,			attention_mask=__lowerCamelCase				,			past_key_values=__lowerCamelCase				,			position_ids=__lowerCamelCase				,			)

			__lowercase					=    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,			dtype='i4'      )
			__lowercase					=    model(
			    input_ids[:, -1:]				,			attention_mask=__lowerCamelCase				,			past_key_values=outputs_cache.past_key_values				,			position_ids=__lowerCamelCase				,			)

			__lowercase					=    model(__lowerCamelCase      )

			__lowercase					=    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
			self.parent.assertTrue(diff < 1E-3				,			msg=F"Max diff is {diff}"      )






	def       UpperCAmelCase (	self       :    Dict				,			__lowerCamelCase       :    Optional[int]				,			__lowerCamelCase       :    str				,			__lowerCamelCase       :    List[Any]				,			__lowerCamelCase       :    Tuple      )					->   str:

			'''simple docstring'''




			__lowercase					=    20
			__lowercase					=    model_class_name(__lowerCamelCase      )

			__lowercase					=    jnp.concatenate(
			    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])      )]				,			axis=-1				,			)

			__lowercase					=    model.init_cache(input_ids.shape[0]				,			__lowerCamelCase      )
			__lowercase					=    jnp.broadcast_to(
			    jnp.arange(input_ids.shape[-1] - 1      )[None, :]				,			(input_ids.shape[0], input_ids.shape[-1] - 1)      )

			__lowercase					=    model(
			    input_ids[:, :-1]				,			attention_mask=__lowerCamelCase				,			past_key_values=__lowerCamelCase				,			position_ids=__lowerCamelCase				,			)
			__lowercase					=    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,			dtype='i4'      )
			__lowercase					=    model(
			    input_ids[:, -1:]				,			past_key_values=outputs_cache.past_key_values				,			attention_mask=__lowerCamelCase				,			position_ids=__lowerCamelCase				,			)

			__lowercase					=    model(__lowerCamelCase				,			attention_mask=__lowerCamelCase      )
			__lowercase					=    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
			self.parent.assertTrue(diff < 1E-3				,			msg=F"Max diff is {diff}"      )





@require_sentencepiece
@require_flax
class 		snake_case_      (       UpperCAmelCase_    ,    UpperCAmelCase_    ,    unittest.TestCase	):







	'''simple docstring'''
	__UpperCamelCase        =		(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
	__UpperCamelCase        =		(FlaxXGLMForCausalLM,) if is_flax_available() else ()




	def       UpperCAmelCase (	self       :    Dict      )					->   List[str]:

			'''simple docstring'''




			__lowercase					=    FlaxXGLMModelTester(self      )




	def       UpperCAmelCase (	self       :    List[str]      )					->   List[Any]:

			'''simple docstring'''




			for model_class_name in self.all_model_classes:
					__lowercase	,			__lowercase	,			__lowercase					=    self.model_tester.prepare_config_and_inputs()
					self.model_tester.check_use_cache_forward(__lowerCamelCase				,			__lowerCamelCase				,			__lowerCamelCase				,			__lowerCamelCase      )




	def       UpperCAmelCase (	self       :    Union[str, Any]      )					->   int:

			'''simple docstring'''




			for model_class_name in self.all_model_classes:
					__lowercase	,			__lowercase	,			__lowercase					=    self.model_tester.prepare_config_and_inputs()
					self.model_tester.check_use_cache_forward_with_attn_mask(
					    __lowerCamelCase				,			__lowerCamelCase				,			__lowerCamelCase				,			__lowerCamelCase      )




	@slow
	def       UpperCAmelCase (	self       :    Any      )					->   int:

			'''simple docstring'''




			__lowercase					=    XGLMTokenizer.from_pretrained('XGLM'				,			padding_side='left'      )
			__lowercase					=    tokenizer(['Hello this is a long string', 'Hey']				,			return_tensors='np'				,			padding=__lowerCamelCase				,			truncation=__lowerCamelCase      )

			__lowercase					=    FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M'      )
			__lowercase					=    1
			__lowercase					=    False

			__lowercase					=    jax.jit(model.generate      )

			__lowercase					=    jit_generate(inputs['input_ids']				,			attention_mask=inputs['attention_mask']      ).sequences

			__lowercase					=    tokenizer.batch_decode(__lowerCamelCase				,			skip_special_tokens=__lowerCamelCase      )

			__lowercase					=    [
			    'Hello this is a long string of questions, but I\'m not sure if I\'m',
			    'Hey, I\'m a newbie to the forum and I\'',
			]

			self.assertListEqual(__lowerCamelCase				,			__lowerCamelCase      )




	@is_pt_flax_cross_test
	def       UpperCAmelCase (	self       :    Optional[int]      )					->   Optional[Any]:

			'''simple docstring'''




			__lowercase	,			__lowercase					=    self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__      ):
							# prepare inputs
							__lowercase					=    self._prepare_for_class(__lowerCamelCase				,			__lowerCamelCase      )
							__lowercase					=    {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

							# load corresponding PyTorch class
							__lowercase					=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
							__lowercase					=    getattr(__lowerCamelCase				,			__lowerCamelCase      )

							__lowercase	,			__lowercase					=    pt_inputs['input_ids'].shape
							__lowercase					=    np.random.randint(0				,			seq_length - 1				,			size=(batch_size,)      )
							for batch_idx, start_index in enumerate(__lowerCamelCase      ):
									__lowercase					=    0
									__lowercase					=    1
									__lowercase					=    0
									__lowercase					=    1
							__lowercase					=    pt_model_class(__lowerCamelCase      ).eval()
							# Flax models don't use the `use_cache` option and cache is not returned as a default.
							# So we disable `use_cache` here for PyTorch model.
							__lowercase					=    False
							__lowercase					=    model_class(__lowerCamelCase				,			dtype=jnp.floataa      )
							__lowercase					=    convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,			__lowerCamelCase      )
							__lowercase					=    fx_state

							with torch.no_grad():
									__lowercase					=    pt_model(**__lowerCamelCase      ).to_tuple()

							__lowercase					=    fx_model(**__lowerCamelCase      ).to_tuple()
							self.assertEqual(len(__lowerCamelCase      )				,			len(__lowerCamelCase      )				,			'Output lengths differ between Flax and PyTorch'      )
							for fx_output, pt_output in zip(__lowerCamelCase				,			__lowerCamelCase      ):
									self.assert_almost_equals(fx_output[:, -1]				,			pt_output[:, -1].numpy()				,			4E-2      )

							with tempfile.TemporaryDirectory() as tmpdirname:
									pt_model.save_pretrained(__lowerCamelCase      )
									__lowercase					=    model_class.from_pretrained(__lowerCamelCase				,			from_pt=__lowerCamelCase      )

							__lowercase					=    fx_model_loaded(**__lowerCamelCase      ).to_tuple()
							self.assertEqual(
							    len(__lowerCamelCase      )				,			len(__lowerCamelCase      )				,			'Output lengths differ between Flax and PyTorch'      )
							for fx_output_loaded, pt_output in zip(__lowerCamelCase				,			__lowerCamelCase      ):
									self.assert_almost_equals(fx_output_loaded[:, -1]				,			pt_output[:, -1].numpy()				,			4E-2      )




	@is_pt_flax_cross_test
	def       UpperCAmelCase (	self       :    Dict      )					->   Dict:

			'''simple docstring'''




			__lowercase	,			__lowercase					=    self.model_tester.prepare_config_and_inputs_for_common()
			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__      ):
							# prepare inputs
							__lowercase					=    self._prepare_for_class(__lowerCamelCase				,			__lowerCamelCase      )
							__lowercase					=    {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

							# load corresponding PyTorch class
							__lowercase					=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
							__lowercase					=    getattr(__lowerCamelCase				,			__lowerCamelCase      )

							__lowercase					=    pt_model_class(__lowerCamelCase      ).eval()
							__lowercase					=    False
							__lowercase					=    model_class(__lowerCamelCase				,			dtype=jnp.floataa      )

							__lowercase					=    load_flax_weights_in_pytorch_model(__lowerCamelCase				,			fx_model.params      )
							__lowercase	,			__lowercase					=    pt_inputs['input_ids'].shape
							__lowercase					=    np.random.randint(0				,			seq_length - 1				,			size=(batch_size,)      )
							for batch_idx, start_index in enumerate(__lowerCamelCase      ):
									__lowercase					=    0
									__lowercase					=    1
									__lowercase					=    0
									__lowercase					=    1

							# make sure weights are tied in PyTorch
							pt_model.tie_weights()

							with torch.no_grad():
									__lowercase					=    pt_model(**__lowerCamelCase      ).to_tuple()

							__lowercase					=    fx_model(**__lowerCamelCase      ).to_tuple()
							self.assertEqual(len(__lowerCamelCase      )				,			len(__lowerCamelCase      )				,			'Output lengths differ between Flax and PyTorch'      )
							for fx_output, pt_output in zip(__lowerCamelCase				,			__lowerCamelCase      ):
									self.assert_almost_equals(fx_output[:, -1]				,			pt_output[:, -1].numpy()				,			4E-2      )

							with tempfile.TemporaryDirectory() as tmpdirname:
									fx_model.save_pretrained(__lowerCamelCase      )
									__lowercase					=    pt_model_class.from_pretrained(__lowerCamelCase				,			from_flax=__lowerCamelCase      )

							with torch.no_grad():
									__lowercase					=    pt_model_loaded(**__lowerCamelCase      ).to_tuple()

							self.assertEqual(
							    len(__lowerCamelCase      )				,			len(__lowerCamelCase      )				,			'Output lengths differ between Flax and PyTorch'      )
							for fx_output, pt_output in zip(__lowerCamelCase				,			__lowerCamelCase      ):
									self.assert_almost_equals(fx_output[:, -1]				,			pt_output[:, -1].numpy()				,			4E-2      )






	@slow
	def       UpperCAmelCase (	self       :    str      )					->   str:

			'''simple docstring'''




			for model_class_name in self.all_model_classes:
					__lowercase					=    model_class_name.from_pretrained('facebook/xglm-564M'      )
					__lowercase					=    model(np.ones((1, 1)      )      )
					self.assertIsNotNone(__lowerCamelCase      )











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 		snake_case_      (       UpperCAmelCase_    ,    UpperCAmelCase_    ,    unittest.TestCase	):







	'''simple docstring'''
	__UpperCamelCase        =		UNetaDModel
	__UpperCamelCase        =		'''sample'''




	@property
	def       UpperCAmelCase (	self       :    str      )					->   List[str]:

			'''simple docstring'''




			__lowercase					=    4
			__lowercase					=    14
			__lowercase					=    16

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

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




	@property
	def       UpperCAmelCase (	self       :    List[str]      )					->   Union[str, Any]:

			'''simple docstring'''




			return (4, 14, 16)




	@property
	def       UpperCAmelCase (	self       :    Dict      )					->   List[Any]:

			'''simple docstring'''




			return (4, 14, 16)




	def       UpperCAmelCase (	self       :    Optional[int]      )					->   Dict:

			'''simple docstring'''




			pass




	def       UpperCAmelCase (	self       :    Union[str, Any]      )					->   Dict:

			'''simple docstring'''




			pass




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    List[Any]      )					->   Optional[Any]:

			'''simple docstring'''




			super().test_determinism()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    List[str]      )					->   Dict:

			'''simple docstring'''




			super().test_outputs_equivalence()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    Optional[Any]      )					->   List[Any]:

			'''simple docstring'''




			super().test_from_save_pretrained()




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

			'''simple docstring'''




			super().test_from_save_pretrained_variant()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    Optional[Any]      )					->   int:

			'''simple docstring'''




			super().test_model_from_pretrained()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    Optional[int]      )					->   Tuple:

			'''simple docstring'''




			super().test_output()




	def       UpperCAmelCase (	self       :    List[Any]      )					->   List[str]:

			'''simple docstring'''




			__lowercase					=    {
			    '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',
			}
			__lowercase					=    self.dummy_input
			return init_dict, inputs_dict




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    Optional[int]      )					->   Any:

			'''simple docstring'''




			__lowercase	,			__lowercase					=    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      )
			__lowercase					=    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       UpperCAmelCase (	self       :    List[str]      )					->   Dict:

			'''simple docstring'''




			__lowercase					=    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      )

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

			with torch.no_grad():
					__lowercase					=    model(__lowerCamelCase				,			__lowerCamelCase      ).sample.permute(0				,			2				,			1      )

			__lowercase					=    output[0, -3:, -3:].flatten()
			# fmt: off
			__lowercase					=    torch.tensor([-2.137172, 1.1426016, 0.3688687, -0.766922, 0.7303146, 0.11038864, -0.4760633, 0.13270172, 0.02591348]      )
			# fmt: on
			self.assertTrue(torch.allclose(__lowerCamelCase				,			__lowerCamelCase				,			rtol=1E-3      )      )




	def       UpperCAmelCase (	self       :    List[Any]      )					->   int:

			'''simple docstring'''




			pass






	@slow
	def       UpperCAmelCase (	self       :    List[str]      )					->   Dict:

			'''simple docstring'''




			__lowercase					=    'harmonai/maestro-150k'
			__lowercase					=    UNetaDModel.from_pretrained(__lowerCamelCase				,			subfolder='unet'      )
			model.to(__lowerCamelCase      )

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

			with torch.no_grad():
					__lowercase					=    model(__lowerCamelCase				,			__lowerCamelCase      ).sample

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

			assert (output_sum - 224.0896).abs() < 0.5
			assert (output_max - 0.0607).abs() < 4E-4





class 		snake_case_      (       UpperCAmelCase_    ,    UpperCAmelCase_    ,    unittest.TestCase	):







	'''simple docstring'''
	__UpperCamelCase        =		UNetaDModel
	__UpperCamelCase        =		'''sample'''




	@property
	def       UpperCAmelCase (	self       :    int      )					->   Union[str, Any]:

			'''simple docstring'''




			__lowercase					=    4
			__lowercase					=    14
			__lowercase					=    16

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

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




	@property
	def       UpperCAmelCase (	self       :    List[str]      )					->   Union[str, Any]:

			'''simple docstring'''




			return (4, 14, 16)




	@property
	def       UpperCAmelCase (	self       :    Tuple      )					->   Union[str, Any]:

			'''simple docstring'''




			return (4, 14, 1)




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    Dict      )					->   str:

			'''simple docstring'''




			super().test_determinism()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    str      )					->   int:

			'''simple docstring'''




			super().test_outputs_equivalence()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    str      )					->   List[Any]:

			'''simple docstring'''




			super().test_from_save_pretrained()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    List[Any]      )					->   int:

			'''simple docstring'''




			super().test_from_save_pretrained_variant()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    List[str]      )					->   Dict:

			'''simple docstring'''




			super().test_model_from_pretrained()




	@unittest.skipIf(torch_device == 'mps'				,			'mish op not supported in MPS'      )
	def       UpperCAmelCase (	self       :    List[str]      )					->   Tuple:

			'''simple docstring'''




			__lowercase	,			__lowercase					=    self.prepare_init_args_and_inputs_for_common()
			__lowercase					=    self.model_class(**__lowerCamelCase      )
			model.to(__lowerCamelCase      )
			model.eval()

			with torch.no_grad():
					__lowercase					=    model(**__lowerCamelCase      )

					if isinstance(__lowerCamelCase				,			__lowerCamelCase      ):
							__lowercase					=    output.sample

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




	def       UpperCAmelCase (	self       :    Optional[int]      )					->   str:

			'''simple docstring'''




			pass




	def       UpperCAmelCase (	self       :    Optional[int]      )					->   Optional[int]:

			'''simple docstring'''




			pass




	def       UpperCAmelCase (	self       :    Optional[Any]      )					->   Optional[Any]:

			'''simple docstring'''




			__lowercase					=    {
			    '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',
			}
			__lowercase					=    self.dummy_input
			return init_dict, inputs_dict




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

			'''simple docstring'''




			__lowercase	,			__lowercase					=    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      )
			__lowercase					=    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       UpperCAmelCase (	self       :    Union[str, Any]      )					->   Optional[Any]:

			'''simple docstring'''




			__lowercase	,			__lowercase					=    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      )

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

			with torch.no_grad():
					__lowercase					=    value_function(__lowerCamelCase				,			__lowerCamelCase      ).sample

			# fmt: off
			__lowercase					=    torch.tensor([165.25] * seq_len      )
			# fmt: on
			self.assertTrue(torch.allclose(__lowerCamelCase				,			__lowerCamelCase				,			rtol=1E-3      )      )






	def       UpperCAmelCase (	self       :    Dict      )					->   str:

			'''simple docstring'''




			pass




