


'''simple docstring'''


from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase__     						=  logging.get_logger(__name__)

lowerCamelCase__     						=  {
    'google/reformer-crime-and-punishment': (
        'https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/config.json'
    ),
    'google/reformer-enwik8': 'https://huggingface.co/google/reformer-enwik8/resolve/main/config.json',
}






class 						lowerCAmelCase__      (      UpperCAmelCase__						):
			lowerCAmelCase			:						List[Any]							 =  "reformer"
			lowerCAmelCase			:						Optional[int]							 =  ["past_buckets_states"]
			lowerCAmelCase			:						int							 =  {}




			def __init__(		self	:  int						,	lowerCamelCase__	:  Optional[int]=64						,	lowerCamelCase__	:  int=["local", "lsh", "local", "lsh", "local", "lsh"]						,	lowerCamelCase__	:  Any=1.0						,	lowerCamelCase__	:  List[str]=True						,	lowerCamelCase__	:  Dict=[64, 64]						,	lowerCamelCase__	:  Tuple=[64, 1_92]						,	lowerCamelCase__	:  Union[str, Any]=0						,	lowerCamelCase__	:  Optional[int]=2						,	lowerCamelCase__	:  Dict=5_12						,	lowerCamelCase__	:  Union[str, Any]=None						,	lowerCamelCase__	:  Tuple="relu"						,	lowerCamelCase__	:  List[str]=0.0_5						,	lowerCamelCase__	:  List[str]=2_56						,	lowerCamelCase__	:  Tuple=0.0_2						,	lowerCamelCase__	:  Tuple=False						,	lowerCamelCase__	:  int=1E-12						,	lowerCamelCase__	:  Union[str, Any]=1						,	lowerCamelCase__	:  Any=0						,	lowerCamelCase__	:  List[str]=0.0_5						,	lowerCamelCase__	:  Optional[int]=64						,	lowerCamelCase__	:  Tuple=64						,	lowerCamelCase__	:  Union[str, Any]=0.0						,	lowerCamelCase__	:  Optional[Any]=1						,	lowerCamelCase__	:  List[str]=0						,	lowerCamelCase__	:  Optional[Any]=40_96						,	lowerCamelCase__	:  str=12						,	lowerCamelCase__	:  Any=None						,	lowerCamelCase__	:  Tuple=1						,	lowerCamelCase__	:  List[str]=0						,	lowerCamelCase__	:  Optional[Any]=3_20						,	lowerCamelCase__	:  List[Any]=False						,	lowerCamelCase__	:  Tuple=True						,	lowerCamelCase__	:  str=None						,	**lowerCamelCase__	:  int						,	) ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   hash_seed
					_UpperCAmelCase		:							int             =   vocab_size
					_UpperCAmelCase		:							List[str]             =   attention_head_size
					_UpperCAmelCase		:							int             =   hidden_size
					_UpperCAmelCase		:							List[str]             =   num_attention_heads
					_UpperCAmelCase		:							Optional[Any]             =   num_hashes
					_UpperCAmelCase		:							Optional[int]             =   len(lowerCamelCase__			)
					_UpperCAmelCase		:							Any             =   tuple(lowerCamelCase__			) if isinstance(lowerCamelCase__						,	lowerCamelCase__			) else num_buckets
					_UpperCAmelCase		:							int             =   lsh_attn_chunk_length
					_UpperCAmelCase		:							str             =   local_attn_chunk_length
					_UpperCAmelCase		:							List[str]             =   lsh_num_chunks_after
					_UpperCAmelCase		:							List[str]             =   lsh_num_chunks_before
					_UpperCAmelCase		:							str             =   local_num_chunks_after
					_UpperCAmelCase		:							Dict             =   local_num_chunks_before
					_UpperCAmelCase		:							Optional[Any]             =   hidden_act
					_UpperCAmelCase		:							Optional[Any]             =   feed_forward_size
					_UpperCAmelCase		:							str             =   hidden_dropout_prob
					_UpperCAmelCase		:							List[Any]             =   lsh_attention_probs_dropout_prob
					_UpperCAmelCase		:							Tuple             =   local_attention_probs_dropout_prob
					_UpperCAmelCase		:							Any             =   max_position_embeddings
					_UpperCAmelCase		:							str             =   initializer_range
					_UpperCAmelCase		:							Tuple             =   layer_norm_eps
					_UpperCAmelCase		:							Optional[int]             =   axial_pos_embds
					_UpperCAmelCase		:							Any             =   tuple(lowerCamelCase__			)
					_UpperCAmelCase		:							Union[str, Any]             =   tuple(lowerCamelCase__			)
					_UpperCAmelCase		:							Tuple             =   axial_norm_std
					_UpperCAmelCase		:							Union[str, Any]             =   chunk_size_lm_head
					_UpperCAmelCase		:							Any             =   attn_layers
					_UpperCAmelCase		:							Union[str, Any]             =   use_cache
					_UpperCAmelCase		:							Optional[Any]             =   classifier_dropout
					super().__init__(
					    pad_token_id=lowerCamelCase__						,	eos_token_id=lowerCamelCase__						,	is_decoder=lowerCamelCase__						,	tie_word_embeddings=lowerCamelCase__						,	**lowerCamelCase__						,	)







'''simple docstring'''


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__      (      UpperCAmelCase__     ,							UpperCAmelCase__     ,							unittest.TestCase						):
			lowerCAmelCase			:						Any							 =  UNetaDModel
			lowerCAmelCase			:						List[Any]							 =  "sample"




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



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   4
					_UpperCAmelCase		:							Any             =   14
					_UpperCAmelCase		:							Any             =   16

					_UpperCAmelCase		:							Optional[int]             =   floats_tensor((batch_size, num_features, seq_len)			).to(lowerCamelCase__			)
					_UpperCAmelCase		:							Dict             =   torch.tensor([10] * batch_size			).to(lowerCamelCase__			)

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




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



					'''simple docstring'''







					return (4, 14, 16)




			@property
			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->Tuple:



					'''simple docstring'''







					return (4, 14, 16)




			def 				lowerCAmelCase__     (		self	:  int			) ->Tuple:



					'''simple docstring'''







					pass




			def 				lowerCAmelCase__     (		self	:  Dict			) ->Tuple:



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					super().test_determinism()




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



					'''simple docstring'''







					super().test_outputs_equivalence()




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



					'''simple docstring'''







					super().test_from_save_pretrained()




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



					'''simple docstring'''







					super().test_from_save_pretrained_variant()




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



					'''simple docstring'''







					super().test_model_from_pretrained()




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



					'''simple docstring'''







					super().test_output()




			def 				lowerCAmelCase__     (		self	:  int			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							Any             =   {
					    "block_out_channels": (32, 64, 1_28, 2_56),
					    "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		:							List[Any]             =   self.dummy_input
					return init_dict, inputs_dict




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



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							int             =   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		:							Dict             =   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 				lowerCAmelCase__     (		self	:  List[str]			) ->Optional[int]:



					'''simple docstring'''







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

					with torch.no_grad():
							_UpperCAmelCase		:							List[Any]             =   model(lowerCamelCase__						,	lowerCamelCase__			).sample.permute(0						,	2						,	1			)

					_UpperCAmelCase		:							Optional[int]             =   output[0, -3:, -3:].flatten()
					# fmt: off
					_UpperCAmelCase		:							Dict             =   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 				lowerCAmelCase__     (		self	:  Any			) ->List[str]:



					'''simple docstring'''







					pass






			@slow
			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   "harmonai/maestro-150k"
					_UpperCAmelCase		:							Tuple             =   UNetaDModel.from_pretrained(lowerCamelCase__						,	subfolder="unet"			)
					model.to(lowerCamelCase__			)

					_UpperCAmelCase		:							List[str]             =   6_55_36
					_UpperCAmelCase		:							Any             =   torch.sin(torch.arange(lowerCamelCase__			)[None, None, :].repeat(1						,	2						,	1			)			).to(lowerCamelCase__			)
					_UpperCAmelCase		:							List[Any]             =   torch.tensor([1]			).to(lowerCamelCase__			)

					with torch.no_grad():
							_UpperCAmelCase		:							List[str]             =   model(lowerCamelCase__						,	lowerCamelCase__			).sample

					_UpperCAmelCase		:							Union[str, Any]             =   output.abs().sum()
					_UpperCAmelCase		:							Any             =   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 						lowerCAmelCase__      (      UpperCAmelCase__     ,							UpperCAmelCase__     ,							unittest.TestCase						):
			lowerCAmelCase			:						Union[str, Any]							 =  UNetaDModel
			lowerCAmelCase			:						Union[str, Any]							 =  "sample"




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



					'''simple docstring'''







					_UpperCAmelCase		:							Tuple             =   4
					_UpperCAmelCase		:							Optional[int]             =   14
					_UpperCAmelCase		:							Optional[int]             =   16

					_UpperCAmelCase		:							Optional[int]             =   floats_tensor((batch_size, num_features, seq_len)			).to(lowerCamelCase__			)
					_UpperCAmelCase		:							Union[str, Any]             =   torch.tensor([10] * batch_size			).to(lowerCamelCase__			)

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




			@property
			def 				lowerCAmelCase__     (		self	:  Dict			) ->Optional[int]:



					'''simple docstring'''







					return (4, 14, 16)




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



					'''simple docstring'''







					return (4, 14, 1)




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



					'''simple docstring'''







					super().test_determinism()




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



					'''simple docstring'''







					super().test_outputs_equivalence()




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



					'''simple docstring'''







					super().test_from_save_pretrained()




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



					'''simple docstring'''







					super().test_from_save_pretrained_variant()




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



					'''simple docstring'''







					super().test_model_from_pretrained()




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



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							str             =   self.prepare_init_args_and_inputs_for_common()
					_UpperCAmelCase		:							Optional[int]             =   self.model_class(**lowerCamelCase__			)
					model.to(lowerCamelCase__			)
					model.eval()

					with torch.no_grad():
							_UpperCAmelCase		:							int             =   model(**lowerCamelCase__			)

							if isinstance(lowerCamelCase__						,	lowerCamelCase__			):
									_UpperCAmelCase		:							List[Any]             =   output.sample

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




			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Dict:



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   {
					    "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, 1_28, 2_56],
					    "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		:							List[Any]             =   self.dummy_input
					return init_dict, inputs_dict




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



					'''simple docstring'''







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



					'''simple docstring'''







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

					with torch.no_grad():
							_UpperCAmelCase		:							int             =   value_function(lowerCamelCase__						,	lowerCamelCase__			).sample

					# fmt: off
					_UpperCAmelCase		:							Union[str, Any]             =   torch.tensor([1_6_5.2_5] * seq_len			)
					# fmt: on
					self.assertTrue(torch.allclose(lowerCamelCase__						,	lowerCamelCase__						,	rtol=1E-3			)			)






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



					'''simple docstring'''







					pass



