




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        __lowerCAmelCase					:







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




									"""simple docstring"""


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







			def UpperCamelCase     (			self       :     Optional[Any] ):




									"""simple docstring"""


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

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

									_UpperCAmelCase						=	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=__snake_case       ,       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       :     Tuple ):




									"""simple docstring"""


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







			def UpperCamelCase     (			self       :     Dict ):




									"""simple docstring"""


									_UpperCAmelCase						=	self.prepare_config_and_inputs()

									_UpperCAmelCase						=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
									_UpperCAmelCase						=	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       :     Any       ,       snake_case__       :     List[Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     Union[str, Any] ):




									"""simple docstring"""


									_UpperCAmelCase						=	20
									_UpperCAmelCase						=	model_class_name(__snake_case )

									_UpperCAmelCase						=	model.init_cache(input_ids.shape[0]       ,       __snake_case )
									_UpperCAmelCase						=	jnp.ones((input_ids.shape[0], max_decoder_length)       ,       dtype="i4" )

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

									_UpperCAmelCase						=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,       dtype="i4" )
									_UpperCAmelCase						=	model(
									    input_ids[:, -1:]       ,       attention_mask=__snake_case       ,       past_key_values=outputs_cache.past_key_values       ,       position_ids=__snake_case       ,       )

									_UpperCAmelCase						=	model(__snake_case )

									_UpperCAmelCase						=	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       :     Any       ,       snake_case__       :     str       ,       snake_case__       :     Any       ,       snake_case__       :     int       ,       snake_case__       :     List[str] ):




									"""simple docstring"""


									_UpperCAmelCase						=	20
									_UpperCAmelCase						=	model_class_name(__snake_case )

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

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

									_UpperCAmelCase						=	model(
									    input_ids[:, :-1]       ,       attention_mask=__snake_case       ,       past_key_values=__snake_case       ,       position_ids=__snake_case       ,       )
									_UpperCAmelCase						=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,       dtype="i4" )
									_UpperCAmelCase						=	model(
									    input_ids[:, -1:]       ,       past_key_values=outputs_cache.past_key_values       ,       attention_mask=__snake_case       ,       position_ids=__snake_case       ,       )

									_UpperCAmelCase						=	model(__snake_case       ,       attention_mask=__snake_case )
									_UpperCAmelCase						=	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        __lowerCAmelCase					(			UpperCAmelCase__   ,						UpperCAmelCase__   ,						unittest.TestCase					):
			snake_case_ :					Dict						=	(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
			snake_case_ :					Tuple						=	(FlaxXGLMForCausalLM,) if is_flax_available() else ()







			def UpperCamelCase     (			self       :     int ):




									"""simple docstring"""


									_UpperCAmelCase						=	FlaxXGLMModelTester(self )







			def UpperCamelCase     (			self       :     str ):




									"""simple docstring"""


									for model_class_name in self.all_model_classes:
															_UpperCAmelCase						=	self.model_tester.prepare_config_and_inputs()
															self.model_tester.check_use_cache_forward(__snake_case       ,       __snake_case       ,       __snake_case       ,       __snake_case )







			def UpperCamelCase     (			self       :     Union[str, Any] ):




									"""simple docstring"""


									for model_class_name in self.all_model_classes:
															_UpperCAmelCase						=	self.model_tester.prepare_config_and_inputs()
															self.model_tester.check_use_cache_forward_with_attn_mask(
															    __snake_case       ,       __snake_case       ,       __snake_case       ,       __snake_case )







			@slow
			def UpperCamelCase     (			self       :     int ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	FlaxXGLMForCausalLM.from_pretrained("facebook/xglm-564M" )
									_UpperCAmelCase						=	1
									_UpperCAmelCase						=	False

									_UpperCAmelCase						=	jax.jit(model.generate )

									_UpperCAmelCase						=	jit_generate(inputs["input_ids"]       ,       attention_mask=inputs["attention_mask"] ).sequences

									_UpperCAmelCase						=	tokenizer.batch_decode(__snake_case       ,       skip_special_tokens=__snake_case )

									_UpperCAmelCase						=	[
									    "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(__snake_case       ,       __snake_case )







			@is_pt_flax_cross_test
			def UpperCamelCase     (			self       :     Tuple ):




									"""simple docstring"""


									_UpperCAmelCase						=	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
																					_UpperCAmelCase						=	self._prepare_for_class(__snake_case       ,       __snake_case )
																					_UpperCAmelCase						=	{k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

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

																					_UpperCAmelCase						=	pt_inputs["input_ids"].shape
																					_UpperCAmelCase						=	np.random.randint(0       ,       seq_length - 1       ,       size=(batch_size,) )
																					for batch_idx, start_index in enumerate(__snake_case ):
																											_UpperCAmelCase						=	0
																											_UpperCAmelCase						=	1
																											_UpperCAmelCase						=	0
																											_UpperCAmelCase						=	1
																					_UpperCAmelCase						=	pt_model_class(__snake_case ).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.
																					_UpperCAmelCase						=	False
																					_UpperCAmelCase						=	model_class(__snake_case       ,       dtype=jnp.floataa )
																					_UpperCAmelCase						=	convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,       __snake_case )
																					_UpperCAmelCase						=	fx_state

																					with torch.no_grad():
																											_UpperCAmelCase						=	pt_model(**__snake_case ).to_tuple()

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

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

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







			@is_pt_flax_cross_test
			def UpperCamelCase     (			self       :     Dict ):




									"""simple docstring"""


									_UpperCAmelCase						=	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
																					_UpperCAmelCase						=	self._prepare_for_class(__snake_case       ,       __snake_case )
																					_UpperCAmelCase						=	{k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

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

																					_UpperCAmelCase						=	pt_model_class(__snake_case ).eval()
																					_UpperCAmelCase						=	False
																					_UpperCAmelCase						=	model_class(__snake_case       ,       dtype=jnp.floataa )

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

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

																					with torch.no_grad():
																											_UpperCAmelCase						=	pt_model(**__snake_case ).to_tuple()

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

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

																					with torch.no_grad():
																											_UpperCAmelCase						=	pt_model_loaded(**__snake_case ).to_tuple()

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







			@slow
			def UpperCamelCase     (			self       :     Any ):




									"""simple docstring"""


									for model_class_name in self.all_model_classes:
															_UpperCAmelCase						=	model_class_name.from_pretrained("facebook/xglm-564M" )
															_UpperCAmelCase						=	model(np.ones((1, 1) ) )
															self.assertIsNotNone(__snake_case )






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 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				int      =							UNetaDModel
							lowerCamelCase__:				Any      =							'''sample'''



							@property
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     4
								UpperCAmelCase:					Optional[Any] 						=     1_4
								UpperCAmelCase:					str 						=     1_6

								UpperCAmelCase:					Optional[int] 						=     floats_tensor((batch_size, num_features, seq_len)							).to(__snake_case							)
								UpperCAmelCase:					Any 						=     torch.tensor([1_0] * batch_size							).to(__snake_case							)

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



							@property
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								return (4, 1_4, 1_6)



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






								"""simple docstring"""



								return (4, 1_4, 1_6)



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								pass



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






								"""simple docstring"""



								pass



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






								"""simple docstring"""



								super().test_determinism()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								super().test_outputs_equivalence()



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






								"""simple docstring"""



								super().test_from_save_pretrained()



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






								"""simple docstring"""



								super().test_from_save_pretrained_variant()



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






								"""simple docstring"""



								super().test_model_from_pretrained()



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






								"""simple docstring"""



								super().test_output()



							def     A__ (     self							)						->     List[str]:






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     {
								    "block_out_channels": (3_2, 6_4, 1_2_8, 2_5_6),
								    "in_channels": 1_4,
								    "out_channels": 1_4,
								    "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:					Tuple 						=     self.dummy_input
								return init_dict, inputs_dict



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Any 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="unet"							)
								self.assertIsNotNone(__snake_case							)
								self.assertEqual(len(loading_info["missing_keys"]							)	,				0							)

								model.to(__snake_case							)
								UpperCAmelCase:					Tuple 						=     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     A__ (     self							)						->     List[str]:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     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:					int 						=     model.config.in_channels
								UpperCAmelCase:					List[str] 						=     1_6
								UpperCAmelCase:					Union[str, Any] 						=     torch.randn((1, seq_len, num_features)							).permute(
								    0	,				2	,				1							)  # match original, we can update values and remove
								UpperCAmelCase:					List[str] 						=     torch.full((num_features,)	,				0							)

								with torch.no_grad():
									UpperCAmelCase:					Union[str, Any] 						=     model(__snake_case	,				__snake_case							).sample.permute(0	,				2	,				1							)

								UpperCAmelCase:					str 						=     output[0, -3:, -3:].flatten()
								# fmt: off
								UpperCAmelCase:					Dict 						=     torch.tensor([-2.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]							)
								# fmt: on
								self.assertTrue(torch.allclose(__snake_case	,				__snake_case	,				rtol=1e-3							)							)



							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								pass




							@slow
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					List[str] 						=     "harmonai/maestro-150k"
								UpperCAmelCase:					Optional[Any] 						=     UNetaDModel.from_pretrained(__snake_case	,				subfolder="unet"							)
								model.to(__snake_case							)

								UpperCAmelCase:					Any 						=     6_5_5_3_6
								UpperCAmelCase:					str 						=     torch.sin(torch.arange(__snake_case							)[None, None, :].repeat(1	,				2	,				1							)							).to(__snake_case							)
								UpperCAmelCase:					Dict 						=     torch.tensor([1]							).to(__snake_case							)

								with torch.no_grad():
									UpperCAmelCase:					Optional[Any] 						=     model(__snake_case	,				__snake_case							).sample

								UpperCAmelCase:					List[Any] 						=     output.abs().sum()
								UpperCAmelCase:					Dict 						=     output.abs().max()

								assert (output_sum - 2_24.08_96).abs() < 0.5
								assert (output_max - 0.06_07).abs() < 4e-4





class 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				Dict      =							UNetaDModel
							lowerCamelCase__:				int      =							'''sample'''



							@property
							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					List[Any] 						=     4
								UpperCAmelCase:					int 						=     1_4
								UpperCAmelCase:					Optional[int] 						=     1_6

								UpperCAmelCase:					List[str] 						=     floats_tensor((batch_size, num_features, seq_len)							).to(__snake_case							)
								UpperCAmelCase:					Any 						=     torch.tensor([1_0] * batch_size							).to(__snake_case							)

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



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






								"""simple docstring"""



								return (4, 1_4, 1_6)



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






								"""simple docstring"""



								return (4, 1_4, 1)



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






								"""simple docstring"""



								super().test_determinism()



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






								"""simple docstring"""



								super().test_outputs_equivalence()



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






								"""simple docstring"""



								super().test_from_save_pretrained()



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






								"""simple docstring"""



								super().test_from_save_pretrained_variant()



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






								"""simple docstring"""



								super().test_model_from_pretrained()



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






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Optional[int] 						=     self.prepare_init_args_and_inputs_for_common()
								UpperCAmelCase:					Tuple 						=     self.model_class(**__snake_case							)
								model.to(__snake_case							)
								model.eval()

								with torch.no_grad():
									UpperCAmelCase:					Union[str, Any] 						=     model(**__snake_case							)

									if isinstance(__snake_case	,				__snake_case							):
										UpperCAmelCase:					List[str] 						=     output.sample

								self.assertIsNotNone(__snake_case							)
								UpperCAmelCase:					Optional[Any] 						=     torch.Size((inputs_dict["sample"].shape[0], 1)							)
								self.assertEqual(output.shape	,				__snake_case	,				"Input and output shapes do not match"							)



							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								pass



							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								pass



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     {
								    "in_channels": 1_4,
								    "out_channels": 1_4,
								    "down_block_types": ["DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"],
								    "up_block_types": [],
								    "out_block_type": "ValueFunction",
								    "mid_block_type": "ValueFunctionMidBlock1D",
								    "block_out_channels": [3_2, 6_4, 1_2_8, 2_5_6],
								    "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:					Union[str, Any] 						=     self.dummy_input
								return init_dict, inputs_dict



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Optional[Any] 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="value_function"							)
								self.assertIsNotNone(__snake_case							)
								self.assertEqual(len(vf_loading_info["missing_keys"]							)	,				0							)

								value_function.to(__snake_case							)
								UpperCAmelCase:					Union[str, Any] 						=     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     A__ (     self							)						->     Any:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Dict 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="value_function"							)
								torch.manual_seed(0							)
								if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

								UpperCAmelCase:					str 						=     value_function.config.in_channels
								UpperCAmelCase:					int 						=     1_4
								UpperCAmelCase:					int 						=     torch.randn((1, seq_len, num_features)							).permute(
								    0	,				2	,				1							)  # match original, we can update values and remove
								UpperCAmelCase:					Union[str, Any] 						=     torch.full((num_features,)	,				0							)

								with torch.no_grad():
									UpperCAmelCase:					Dict 						=     value_function(__snake_case	,				__snake_case							).sample

								# fmt: off
								UpperCAmelCase:					Tuple 						=     torch.tensor([1_65.25] * seq_len							)
								# fmt: on
								self.assertTrue(torch.allclose(__snake_case	,				__snake_case	,				rtol=1e-3							)							)




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






								"""simple docstring"""



								pass
