


"""simple docstring"""



from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
	import accelerate






def       lowerCamelCase					(a_						:List[Any])     ->							Dict:

							if not is_accelerate_available():
														return method
							lowercase		:Optional[int]       							=				version.parse(accelerate.__version__).base_version
							if version.parse(UpperCAmelCase__) < version.parse('''0.17.0'''):
														return method

							def wrapper(self						:List[Any]				,				*a_						:int				,				**a_						:int):
														if hasattr(self				,				'''_hf_hook''') and hasattr(self._hf_hook				,				'''pre_forward'''):
																					self._hf_hook.pre_forward(self)
														return method(self				,				*UpperCAmelCase__				,				**UpperCAmelCase__)

							return wrapper





"""simple docstring"""



import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
 import jax
@require_flax
class 					__magic_name__		(			__UpperCAmelCase			,						unittest.TestCase     ):
       __A			:      str        =   FlaxAutoencoderKL






       @property
       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase		:Dict       							=				4
              lowercase		:Optional[int]       							=				3
              lowercase		:Any       							=				(3_2, 3_2)

              lowercase		:Optional[int]       							=				jax.random.PRNGKey(0				)
              lowercase		:Any       							=				jax.random.uniform(snake_case__  ,     ((batch_size, num_channels) + sizes)				)

              return {"sample": image, "prng_key": prng_key}







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

              '''simple docstring'''






              lowercase		:List[Any]       							=				{
                  '''block_out_channels''': [3_2, 6_4],
                  '''in_channels''': 3,
                  '''out_channels''': 3,
                  '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
                  '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
                  '''latent_channels''': 4,
              }
              lowercase		:Optional[int]       							=				self.dummy_input
              return init_dict, inputs_dict

