



import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
							import jax.numpy as jnp

if is_tf_available():
							import tensorflow as tf

if is_torch_available():
							import torch

class 	lowerCamelCase_    ( unittest.TestCase						):






						'''simple docstring'''







						def 			SCREAMING_SNAKE_CASE__		(	self      :			int )      ->				Dict:
									A							:					str              =						{
									    '''task_specific_params''': {
									        '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 1_28, '''min_length''': 12, '''num_beams''': 4},
									        '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 1_42, '''min_length''': 56, '''num_beams''': 4},
									        '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 62, '''min_length''': 11, '''num_beams''': 6},
									    }
									}
									A							:					List[str]              =						{
									    '''task_specific_params.summarization.length_penalty''': 1.0,
									    '''task_specific_params.summarization.max_length''': 1_28,
									    '''task_specific_params.summarization.min_length''': 12,
									    '''task_specific_params.summarization.num_beams''': 4,
									    '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
									    '''task_specific_params.summarization_cnn.max_length''': 1_42,
									    '''task_specific_params.summarization_cnn.min_length''': 56,
									    '''task_specific_params.summarization_cnn.num_beams''': 4,
									    '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
									    '''task_specific_params.summarization_xsum.max_length''': 62,
									    '''task_specific_params.summarization_xsum.min_length''': 11,
									    '''task_specific_params.summarization_xsum.num_beams''': 6,
									}

									self.assertEqual(flatten_dict(UpperCamelCase__ )	,	UpperCamelCase__ )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				Union[str, Any]:
									A							:					Tuple              =						np.random.randn(3	,	4 )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__ )	,	x.transpose() ) )

									A							:					str              =						np.random.randn(3	,	4	,	5 )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__	,	axes=(1, 2, 0) )	,	x.transpose((1, 2, 0) ) ) )







						@require_torch
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Union[str, Any] )      ->				Optional[Any]:
									A							:					Tuple              =						np.random.randn(3	,	4 )
									A							:					str              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__ )	,	transpose(UpperCamelCase__ ).numpy() ) )

									A							:					Optional[int]              =						np.random.randn(3	,	4	,	5 )
									A							:					Optional[int]              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__	,	axes=(1, 2, 0) )	,	transpose(UpperCamelCase__	,	axes=(1, 2, 0) ).numpy() ) )







						@require_tf
						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				Optional[int]:
									A							:					List[Any]              =						np.random.randn(3	,	4 )
									A							:					Union[str, Any]              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__ )	,	transpose(UpperCamelCase__ ).numpy() ) )

									A							:					Any              =						np.random.randn(3	,	4	,	5 )
									A							:					Dict              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__	,	axes=(1, 2, 0) )	,	transpose(UpperCamelCase__	,	axes=(1, 2, 0) ).numpy() ) )







						@require_flax
						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[str] )      ->				Union[str, Any]:
									A							:					List[str]              =						np.random.randn(3	,	4 )
									A							:					List[str]              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__ )	,	np.asarray(transpose(UpperCamelCase__ ) ) ) )

									A							:					Any              =						np.random.randn(3	,	4	,	5 )
									A							:					Any              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(transpose(UpperCamelCase__	,	axes=(1, 2, 0) )	,	np.asarray(transpose(UpperCamelCase__	,	axes=(1, 2, 0) ) ) ) )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Any )      ->				Optional[Any]:
									A							:					Dict              =						np.random.randn(3	,	4 )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(4, 3) )	,	np.reshape(UpperCamelCase__	,	(4, 3) ) ) )

									A							:					Union[str, Any]              =						np.random.randn(3	,	4	,	5 )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(12, 5) )	,	np.reshape(UpperCamelCase__	,	(12, 5) ) ) )







						@require_torch
						def 			SCREAMING_SNAKE_CASE__		(	self      :			str )      ->				List[Any]:
									A							:					Dict              =						np.random.randn(3	,	4 )
									A							:					Union[str, Any]              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(4, 3) )	,	reshape(UpperCamelCase__	,	(4, 3) ).numpy() ) )

									A							:					int              =						np.random.randn(3	,	4	,	5 )
									A							:					Optional[int]              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(12, 5) )	,	reshape(UpperCamelCase__	,	(12, 5) ).numpy() ) )







						@require_tf
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[int] )      ->				Optional[Any]:
									A							:					Optional[int]              =						np.random.randn(3	,	4 )
									A							:					Optional[int]              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(4, 3) )	,	reshape(UpperCamelCase__	,	(4, 3) ).numpy() ) )

									A							:					Any              =						np.random.randn(3	,	4	,	5 )
									A							:					Optional[int]              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(12, 5) )	,	reshape(UpperCamelCase__	,	(12, 5) ).numpy() ) )







						@require_flax
						def 			SCREAMING_SNAKE_CASE__		(	self      :			str )      ->				str:
									A							:					Optional[int]              =						np.random.randn(3	,	4 )
									A							:					Union[str, Any]              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(4, 3) )	,	np.asarray(reshape(UpperCamelCase__	,	(4, 3) ) ) ) )

									A							:					int              =						np.random.randn(3	,	4	,	5 )
									A							:					List[str]              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(reshape(UpperCamelCase__	,	(12, 5) )	,	np.asarray(reshape(UpperCamelCase__	,	(12, 5) ) ) ) )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Union[str, Any] )      ->				int:
									A							:					Dict              =						np.random.randn(1	,	3	,	4 )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__ )	,	np.squeeze(UpperCamelCase__ ) ) )

									A							:					Dict              =						np.random.randn(1	,	4	,	1	,	5 )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__	,	axis=2 )	,	np.squeeze(UpperCamelCase__	,	axis=2 ) ) )







						@require_torch
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[int] )      ->				Optional[int]:
									A							:					List[Any]              =						np.random.randn(1	,	3	,	4 )
									A							:					List[str]              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__ )	,	squeeze(UpperCamelCase__ ).numpy() ) )

									A							:					Dict              =						np.random.randn(1	,	4	,	1	,	5 )
									A							:					Optional[int]              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__	,	axis=2 )	,	squeeze(UpperCamelCase__	,	axis=2 ).numpy() ) )







						@require_tf
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[Any] )      ->				Union[str, Any]:
									A							:					int              =						np.random.randn(1	,	3	,	4 )
									A							:					Dict              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__ )	,	squeeze(UpperCamelCase__ ).numpy() ) )

									A							:					List[Any]              =						np.random.randn(1	,	4	,	1	,	5 )
									A							:					Optional[Any]              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__	,	axis=2 )	,	squeeze(UpperCamelCase__	,	axis=2 ).numpy() ) )







						@require_flax
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Dict )      ->				List[Any]:
									A							:					int              =						np.random.randn(1	,	3	,	4 )
									A							:					str              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__ )	,	np.asarray(squeeze(UpperCamelCase__ ) ) ) )

									A							:					str              =						np.random.randn(1	,	4	,	1	,	5 )
									A							:					Optional[int]              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(squeeze(UpperCamelCase__	,	axis=2 )	,	np.asarray(squeeze(UpperCamelCase__	,	axis=2 ) ) ) )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Dict )      ->				List[str]:
									A							:					Optional[Any]              =						np.random.randn(3	,	4 )
									self.assertTrue(np.allclose(expand_dims(UpperCamelCase__	,	axis=1 )	,	np.expand_dims(UpperCamelCase__	,	axis=1 ) ) )







						@require_torch
						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[str] )      ->				Dict:
									A							:					List[Any]              =						np.random.randn(3	,	4 )
									A							:					Optional[Any]              =						torch.tensor(UpperCamelCase__ )
									self.assertTrue(np.allclose(expand_dims(UpperCamelCase__	,	axis=1 )	,	expand_dims(UpperCamelCase__	,	axis=1 ).numpy() ) )







						@require_tf
						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				List[Any]:
									A							:					int              =						np.random.randn(3	,	4 )
									A							:					str              =						tf.constant(UpperCamelCase__ )
									self.assertTrue(np.allclose(expand_dims(UpperCamelCase__	,	axis=1 )	,	expand_dims(UpperCamelCase__	,	axis=1 ).numpy() ) )







						@require_flax
						def 			SCREAMING_SNAKE_CASE__		(	self      :			int )      ->				Optional[Any]:
									A							:					Union[str, Any]              =						np.random.randn(3	,	4 )
									A							:					Optional[int]              =						jnp.array(UpperCamelCase__ )
									self.assertTrue(np.allclose(expand_dims(UpperCamelCase__	,	axis=1 )	,	np.asarray(expand_dims(UpperCamelCase__	,	axis=1 ) ) ) )




import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch

@require_torch
class 	lowerCamelCase_    :






						'''simple docstring'''







						@property
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[int] )      ->				str:
									return self.get_dummy_input()







						@property
						def 			SCREAMING_SNAKE_CASE__		(	self      :			int )      ->				Optional[Any]:
									if self.block_type == "down":
												return (4, 32, 16, 16)
									elif self.block_type == "mid":
												return (4, 32, 32, 32)
									elif self.block_type == "up":
												return (4, 32, 64, 64)

									raise ValueError(F"""'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'.""" )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Dict	,	__lowerCamelCase      :			Union[str, Any]=True	,	__lowerCamelCase      :			int=False	,	__lowerCamelCase      :			int=False	,	__lowerCamelCase      :			Optional[int]=False	,	)      ->				Dict:
									A							:					Optional[Any]              =						4
									A							:					List[str]              =						32
									A							:					Any              =						(32, 32)

									A							:					str              =						torch.manual_seed(0 )
									A							:					int              =						torch.device(__lowerCamelCase )
									A							:					List[str]              =						(batch_size, num_channels) + sizes
									A							:					Dict              =						randn_tensor(__lowerCamelCase	,	generator=__lowerCamelCase	,	device=__lowerCamelCase )
									A							:					int              =						{"hidden_states": hidden_states}

									if include_temb:
												A							:					Any              =						1_28
												A							:					List[str]              =						randn_tensor((batch_size, temb_channels)	,	generator=__lowerCamelCase	,	device=__lowerCamelCase )

									if include_res_hidden_states_tuple:
												A							:					str              =						torch.manual_seed(1 )
												A							:					Tuple              =						(randn_tensor(__lowerCamelCase	,	generator=__lowerCamelCase	,	device=__lowerCamelCase ),)

									if include_encoder_hidden_states:
												A							:					Dict              =						floats_tensor((batch_size, 32, 32) ).to(__lowerCamelCase )

									if include_skip_sample:
												A							:					Optional[int]              =						randn_tensor(((batch_size, 3) + sizes)	,	generator=__lowerCamelCase	,	device=__lowerCamelCase )

									return dummy_input







						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				Union[str, Any]:
									A							:					Dict              =						{
									    "in_channels": 32,
									    "out_channels": 32,
									    "temb_channels": 1_28,
									}
									if self.block_type == "up":
												A							:					Dict              =						32

									if self.block_type == "mid":
												init_dict.pop("out_channels" )

									A							:					str              =						self.dummy_input
									return init_dict, inputs_dict







						def 			SCREAMING_SNAKE_CASE__		(	self      :			str	,	__lowerCamelCase      :			Optional[int] )      ->				Union[str, Any]:
									A       , A							:					str              =						self.prepare_init_args_and_inputs_for_common()
									A							:					List[Any]              =						self.block_class(**__lowerCamelCase )
									unet_block.to(__lowerCamelCase )
									unet_block.eval()

									with torch.no_grad():
												A							:					int              =						unet_block(**__lowerCamelCase )

									if isinstance(__lowerCamelCase	,	__lowerCamelCase ):
												A							:					Union[str, Any]              =						output[0]

									self.assertEqual(output.shape	,	self.output_shape )

									A							:					Any              =						output[0, -1, -3:, -3:]
									A							:					Union[str, Any]              =						torch.tensor(__lowerCamelCase ).to(__lowerCamelCase )
									assert torch_all_close(output_slice.flatten()	,	__lowerCamelCase	,	atol=5e-3 )







						@unittest.skipIf(torch_device == "mps"	,	"Training is not supported in mps" )
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Any )      ->				Dict:
									A       , A							:					Tuple              =						self.prepare_init_args_and_inputs_for_common()
									A							:					str              =						self.block_class(**__lowerCamelCase )
									model.to(__lowerCamelCase )
									model.train()
									A							:					Optional[int]              =						model(**__lowerCamelCase )

									if isinstance(__lowerCamelCase	,	__lowerCamelCase ):
												A							:					Optional[Any]              =						output[0]

									A							:					List[str]              =						torch.device(__lowerCamelCase )
									A							:					List[str]              =						randn_tensor(output.shape	,	device=__lowerCamelCase )
									A							:					Dict              =						torch.nn.functional.mse_loss(__lowerCamelCase	,	__lowerCamelCase )
									loss.backward()