


from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


A	:		Dict    =       logging.get_logger(__name__)

# General docstring
A	:		str    =       'RegNetConfig'

# Base docstring
A	:		Optional[Any]    =       'facebook/regnet-y-040'
A	:		Optional[Any]    =       [1, 1_0_8_8, 7, 7]

# Image classification docstring
A	:		Optional[Any]    =       'facebook/regnet-y-040'
A	:		Any    =       'tabby, tabby cat'

A	:		List[Any]    =       [
    'facebook/regnet-y-040',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]



class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						str				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int = 3				, _UpperCAmelCase   :						int = 1				, _UpperCAmelCase   :						int = 1				, _UpperCAmelCase   :						Optional[str] = "relu"				, **_UpperCAmelCase   :						int				, )						->       Tuple:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							# The padding and conv has been verified in
							# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
							lowercase__					      =		tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2		)
							lowercase__					      =		tf.keras.layers.ConvaD(
							    filters=_UpperCAmelCase				, kernel_size=_UpperCAmelCase				, strides=_UpperCAmelCase				, padding="""VALID"""				, groups=_UpperCAmelCase				, use_bias=_UpperCAmelCase				, name="""convolution"""				, )
							lowercase__					      =		tf.keras.layers.BatchNormalization(epsilon=1E-5				, momentum=0.9				, name="""normalization"""		)
							lowercase__					      =		ACTaFN[activation] if activation is not None else tf.identity


					def   lowerCamelCase__					(self   :						Dict				, _UpperCAmelCase   :						List[Any]		)						->       Optional[Any]:




							"""simple docstring"""


							lowercase__					      =		self.convolution(self.padding(_UpperCAmelCase		)		)
							lowercase__					      =		self.normalization(_UpperCAmelCase		)
							lowercase__					      =		self.activation(_UpperCAmelCase		)
							return hidden_state






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						Tuple				, _UpperCAmelCase   :						RegNetConfig				, **_UpperCAmelCase   :						Optional[Any]		)						->       Union[str, Any]:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		config.num_channels
							lowercase__					      =		TFRegNetConvLayer(
							    out_channels=config.embedding_size				, kernel_size=3				, stride=2				, activation=config.hidden_act				, name="""embedder"""				, )


					def   lowerCamelCase__					(self   :						str				, _UpperCAmelCase   :						List[Any]		)						->       Dict:




							"""simple docstring"""


							lowercase__					      =		shape_list(_UpperCAmelCase		)[1]
							if tf.executing_eagerly() and num_channels != self.num_channels:
									raise ValueError(
									    """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""		)

							# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
							# So change the input format from `NCHW` to `NHWC`.
							# shape = (batch_size, in_height, in_width, in_channels=num_channels)
							lowercase__					      =		tf.transpose(_UpperCAmelCase				, perm=(0, 2, 3, 1)		)
							lowercase__					      =		self.embedder(_UpperCAmelCase		)
							return hidden_state






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						int				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int = 2				, **_UpperCAmelCase   :						Optional[int]		)						->       Dict:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		tf.keras.layers.ConvaD(
							    filters=_UpperCAmelCase				, kernel_size=1				, strides=_UpperCAmelCase				, use_bias=_UpperCAmelCase				, name="""convolution"""		)
							lowercase__					      =		tf.keras.layers.BatchNormalization(epsilon=1E-5				, momentum=0.9				, name="""normalization"""		)


					def   lowerCamelCase__					(self   :						str				, _UpperCAmelCase   :						tf.Tensor				, _UpperCAmelCase   :						bool = False		)						->       tf.Tensor:




							"""simple docstring"""


							return self.normalization(self.convolution(_UpperCAmelCase		)				, training=_UpperCAmelCase		)






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						List[Any]				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int				, **_UpperCAmelCase   :						int		)						->       Dict:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		tf.keras.layers.GlobalAveragePoolingaD(keepdims=_UpperCAmelCase				, name="""pooler"""		)
							lowercase__					      =		[
							    tf.keras.layers.ConvaD(filters=_UpperCAmelCase				, kernel_size=1				, activation="""relu"""				, name="""attention.0"""		),
							    tf.keras.layers.ConvaD(filters=_UpperCAmelCase				, kernel_size=1				, activation="""sigmoid"""				, name="""attention.2"""		),
							]


					def   lowerCamelCase__					(self   :						int				, _UpperCAmelCase   :						List[str]		)						->       Union[str, Any]:




							"""simple docstring"""


							lowercase__					      =		self.pooler(_UpperCAmelCase		)
							for layer_module in self.attention:
									lowercase__					      =		layer_module(_UpperCAmelCase		)
							lowercase__					      =		hidden_state * pooled
							return hidden_state






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						int				, _UpperCAmelCase   :						RegNetConfig				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int = 1				, **_UpperCAmelCase   :						List[Any]		)						->       str:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		in_channels != out_channels or stride != 1
							lowercase__					      =		max(1				, out_channels // config.groups_width		)
							lowercase__					      =		(
							    TFRegNetShortCut(_UpperCAmelCase				, stride=_UpperCAmelCase				, name="""shortcut"""		)
							    if should_apply_shortcut
							    else tf.keras.layers.Activation("""linear"""				, name="""shortcut"""		)
							)
							# `self.layers` instead of `self.layer` because that is a reserved argument.
							lowercase__					      =		[
							    TFRegNetConvLayer(_UpperCAmelCase				, kernel_size=1				, activation=config.hidden_act				, name="""layer.0"""		),
							    TFRegNetConvLayer(
							        _UpperCAmelCase				, stride=_UpperCAmelCase				, groups=_UpperCAmelCase				, activation=config.hidden_act				, name="""layer.1"""		),
							    TFRegNetConvLayer(_UpperCAmelCase				, kernel_size=1				, activation=_UpperCAmelCase				, name="""layer.2"""		),
							]
							lowercase__					      =		ACTaFN[config.hidden_act]


					def   lowerCamelCase__					(self   :						Any				, _UpperCAmelCase   :						str		)						->       List[Any]:




							"""simple docstring"""


							lowercase__					      =		hidden_state
							for layer_module in self.layers:
									lowercase__					      =		layer_module(_UpperCAmelCase		)
							lowercase__					      =		self.shortcut(_UpperCAmelCase		)
							hidden_state += residual
							lowercase__					      =		self.activation(_UpperCAmelCase		)
							return hidden_state






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						Dict				, _UpperCAmelCase   :						RegNetConfig				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int = 1				, **_UpperCAmelCase   :						str		)						->       str:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		in_channels != out_channels or stride != 1
							lowercase__					      =		max(1				, out_channels // config.groups_width		)
							lowercase__					      =		(
							    TFRegNetShortCut(_UpperCAmelCase				, stride=_UpperCAmelCase				, name="""shortcut"""		)
							    if should_apply_shortcut
							    else tf.keras.layers.Activation("""linear"""				, name="""shortcut"""		)
							)
							lowercase__					      =		[
							    TFRegNetConvLayer(_UpperCAmelCase				, kernel_size=1				, activation=config.hidden_act				, name="""layer.0"""		),
							    TFRegNetConvLayer(
							        _UpperCAmelCase				, stride=_UpperCAmelCase				, groups=_UpperCAmelCase				, activation=config.hidden_act				, name="""layer.1"""		),
							    TFRegNetSELayer(_UpperCAmelCase				, reduced_channels=int(round(in_channels / 4		)		)				, name="""layer.2"""		),
							    TFRegNetConvLayer(_UpperCAmelCase				, kernel_size=1				, activation=_UpperCAmelCase				, name="""layer.3"""		),
							]
							lowercase__					      =		ACTaFN[config.hidden_act]


					def   lowerCamelCase__					(self   :						int				, _UpperCAmelCase   :						int		)						->       Dict:




							"""simple docstring"""


							lowercase__					      =		hidden_state
							for layer_module in self.layers:
									lowercase__					      =		layer_module(_UpperCAmelCase		)
							lowercase__					      =		self.shortcut(_UpperCAmelCase		)
							hidden_state += residual
							lowercase__					      =		self.activation(_UpperCAmelCase		)
							return hidden_state






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						List[str]				, _UpperCAmelCase   :						RegNetConfig				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int = 2				, _UpperCAmelCase   :						int = 2				, **_UpperCAmelCase   :						str		)						->       Union[str, Any]:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)

							lowercase__					      =		TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
							lowercase__					      =		[
							    # downsampling is done in the first layer with stride of 2
							    layer(_UpperCAmelCase				, _UpperCAmelCase				, _UpperCAmelCase				, stride=_UpperCAmelCase				, name="""layers.0"""		),
							    *[layer(_UpperCAmelCase				, _UpperCAmelCase				, _UpperCAmelCase				, name=f'''layers.{i+1}'''		) for i in range(depth - 1		)],
							]


					def   lowerCamelCase__					(self   :						Any				, _UpperCAmelCase   :						Any		)						->       Dict:




							"""simple docstring"""


							for layer_module in self.layers:
									lowercase__					      =		layer_module(_UpperCAmelCase		)
							return hidden_state






class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					def __init__(self   :						Union[str, Any]				, _UpperCAmelCase   :						RegNetConfig				, **_UpperCAmelCase   :						Dict		)						->       Union[str, Any]:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		[]
							# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
							self.stages.append(
							    TFRegNetStage(
							        _UpperCAmelCase				, config.embedding_size				, config.hidden_sizes[0]				, stride=2 if config.downsample_in_first_stage else 1				, depth=config.depths[0]				, name="""stages.0"""				, )		)
							lowercase__					      =		zip(config.hidden_sizes				, config.hidden_sizes[1:]		)
							for i, ((in_channels, out_channels), depth) in enumerate(zip(_UpperCAmelCase				, config.depths[1:]		)		):
									self.stages.append(TFRegNetStage(_UpperCAmelCase				, _UpperCAmelCase				, _UpperCAmelCase				, depth=_UpperCAmelCase				, name=f'''stages.{i+1}'''		)		)


					def   lowerCamelCase__					(self   :						List[str]				, _UpperCAmelCase   :						tf.Tensor				, _UpperCAmelCase   :						bool = False				, _UpperCAmelCase   :						bool = True		)						->       TFBaseModelOutputWithNoAttention:




							"""simple docstring"""


							lowercase__					      =		() if output_hidden_states else None

							for stage_module in self.stages:
									if output_hidden_states:
											lowercase__					      =		hidden_states + (hidden_state,)

									lowercase__					      =		stage_module(_UpperCAmelCase		)

							if output_hidden_states:
									lowercase__					      =		hidden_states + (hidden_state,)

							if not return_dict:
									return tuple(v for v in [hidden_state, hidden_states] if v is not None		)

							return TFBaseModelOutputWithNoAttention(last_hidden_state=_UpperCAmelCase				, hidden_states=_UpperCAmelCase		)






@keras_serializable
class 							A   (	tf.keras.layers.Layer							):
					'''simple docstring'''






					A__      =			RegNetConfig

					def __init__(self   :						Optional[Any]				, _UpperCAmelCase   :						str				, **_UpperCAmelCase   :						Any		)						->       List[Any]:




							"""simple docstring"""


							super().__init__(**_UpperCAmelCase		)
							lowercase__					      =		config
							lowercase__					      =		TFRegNetEmbeddings(_UpperCAmelCase				, name="""embedder"""		)
							lowercase__					      =		TFRegNetEncoder(_UpperCAmelCase				, name="""encoder"""		)
							lowercase__					      =		tf.keras.layers.GlobalAveragePoolingaD(keepdims=_UpperCAmelCase				, name="""pooler"""		)


					@unpack_inputs
					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						tf.Tensor				, _UpperCAmelCase   :						Optional[bool] = None				, _UpperCAmelCase   :						Optional[bool] = None				, _UpperCAmelCase   :						bool = False				, )						->       TFBaseModelOutputWithPoolingAndNoAttention:




							"""simple docstring"""


							lowercase__					      =		(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							lowercase__					      =		return_dict if return_dict is not None else self.config.use_return_dict

							lowercase__					      =		self.embedder(_UpperCAmelCase				, training=_UpperCAmelCase		)

							lowercase__					      =		self.encoder(
							    _UpperCAmelCase				, output_hidden_states=_UpperCAmelCase				, return_dict=_UpperCAmelCase				, training=_UpperCAmelCase		)

							lowercase__					      =		encoder_outputs[0]
							lowercase__					      =		self.pooler(_UpperCAmelCase		)

							# Change to NCHW output format have uniformity in the modules
							lowercase__					      =		tf.transpose(_UpperCAmelCase				, perm=(0, 3, 1, 2)		)
							lowercase__					      =		tf.transpose(_UpperCAmelCase				, perm=(0, 3, 1, 2)		)

							# Change the other hidden state outputs to NCHW as well
							if output_hidden_states:
									lowercase__					      =		tuple([tf.transpose(_UpperCAmelCase				, perm=(0, 3, 1, 2)		) for h in encoder_outputs[1]]		)

							if not return_dict:
									return (last_hidden_state, pooled_output) + encoder_outputs[1:]

							return TFBaseModelOutputWithPoolingAndNoAttention(
							    last_hidden_state=_UpperCAmelCase				, pooler_output=_UpperCAmelCase				, hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states				, )






class 							A   (	UpperCAmelCase__							):
					'''simple docstring'''






					A__      =			RegNetConfig
					A__      =			'''regnet'''
					A__      =			'''pixel_values'''

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




							"""simple docstring"""


							return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)				, dtype=tf.floataa		)}






A	:		List[Any]    =       r'\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'

A	:		Optional[Any]    =       r'\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'



@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.'''			,  UpperCAmelCase__			,  )
class 							A   (	UpperCAmelCase__							):
					'''simple docstring'''






					def __init__(self   :						Tuple				, _UpperCAmelCase   :						RegNetConfig				, *_UpperCAmelCase   :						Dict				, **_UpperCAmelCase   :						Any		)						->       Optional[Any]:




							"""simple docstring"""


							super().__init__(_UpperCAmelCase				, *_UpperCAmelCase				, **_UpperCAmelCase		)
							lowercase__					      =		TFRegNetMainLayer(_UpperCAmelCase				, name="""regnet"""		)


					@unpack_inputs
					@add_start_docstrings_to_model_forward(_UpperCAmelCase		)
					@add_code_sample_docstrings(
					    checkpoint=_CHECKPOINT_FOR_DOC				, output_type=_UpperCAmelCase				, config_class=_CONFIG_FOR_DOC				, modality="""vision"""				, expected_output=_EXPECTED_OUTPUT_SHAPE				, )
					def   lowerCamelCase__					(self   :						List[str]				, _UpperCAmelCase   :						tf.Tensor				, _UpperCAmelCase   :						Optional[bool] = None				, _UpperCAmelCase   :						Optional[bool] = None				, _UpperCAmelCase   :						List[Any]=False				, )						->       Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:




							"""simple docstring"""


							lowercase__					      =		(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							lowercase__					      =		return_dict if return_dict is not None else self.config.use_return_dict

							lowercase__					      =		self.regnet(
							    pixel_values=_UpperCAmelCase				, output_hidden_states=_UpperCAmelCase				, return_dict=_UpperCAmelCase				, training=_UpperCAmelCase				, )
							if not return_dict:
									return (outputs[0],) + outputs[1:]

							return TFBaseModelOutputWithPoolingAndNoAttention(
							    last_hidden_state=outputs.last_hidden_state				, pooler_output=outputs.pooler_output				, hidden_states=outputs.hidden_states				, )






@add_start_docstrings(
    '''
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''			,  UpperCAmelCase__			,  )
class 							A   (	UpperCAmelCase__			,  UpperCAmelCase__							):
					'''simple docstring'''






					def __init__(self   :						List[str]				, _UpperCAmelCase   :						RegNetConfig				, *_UpperCAmelCase   :						Any				, **_UpperCAmelCase   :						str		)						->       str:




							"""simple docstring"""


							super().__init__(_UpperCAmelCase				, *_UpperCAmelCase				, **_UpperCAmelCase		)
							lowercase__					      =		config.num_labels
							lowercase__					      =		TFRegNetMainLayer(_UpperCAmelCase				, name="""regnet"""		)
							# classification head
							lowercase__					      =		[
							    tf.keras.layers.Flatten(),
							    tf.keras.layers.Dense(config.num_labels				, name="""classifier.1"""		) if config.num_labels > 0 else tf.identity,
							]


					@unpack_inputs
					@add_start_docstrings_to_model_forward(_UpperCAmelCase		)
					@add_code_sample_docstrings(
					    checkpoint=_IMAGE_CLASS_CHECKPOINT				, output_type=_UpperCAmelCase				, config_class=_CONFIG_FOR_DOC				, expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				, )
					def   lowerCamelCase__					(self   :						Optional[int]				, _UpperCAmelCase   :						tf.Tensor = None				, _UpperCAmelCase   :						tf.Tensor = None				, _UpperCAmelCase   :						bool = None				, _UpperCAmelCase   :						bool = None				, _UpperCAmelCase   :						Optional[int]=False				, )						->       Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:




							"""simple docstring"""


							lowercase__					      =		(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							lowercase__					      =		return_dict if return_dict is not None else self.config.use_return_dict

							lowercase__					      =		self.regnet(
							    _UpperCAmelCase				, output_hidden_states=_UpperCAmelCase				, return_dict=_UpperCAmelCase				, training=_UpperCAmelCase		)

							lowercase__					      =		outputs.pooler_output if return_dict else outputs[1]

							lowercase__					      =		self.classifier[0](_UpperCAmelCase		)
							lowercase__					      =		self.classifier[1](_UpperCAmelCase		)

							lowercase__					      =		None if labels is None else self.hf_compute_loss(labels=_UpperCAmelCase				, logits=_UpperCAmelCase		)

							if not return_dict:
									lowercase__					      =		(logits,) + outputs[2:]
									return ((loss,) + output) if loss is not None else output

							return TFSequenceClassifierOutput(loss=_UpperCAmelCase				, logits=_UpperCAmelCase				, hidden_states=outputs.hidden_states		)





import gc
import unittest

from parameterized import parameterized

from diffusers import FlaxUNetaDConditionModel
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import load_hf_numpy, require_flax, slow


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



@slow
@require_flax
class 							A   (	unittest.TestCase							):
					'''simple docstring'''






					def   lowerCamelCase__					(self   :						Any				, _UpperCAmelCase   :						List[str]				, _UpperCAmelCase   :						Union[str, Any]		)						->       int:




							"""simple docstring"""


							return f'''gaussian_noise_s={seed}_shape={"_".join([str(_UpperCAmelCase		) for s in shape]		)}.npy'''

					def   lowerCamelCase__					(self   :						Any		)						->       int:




							"""simple docstring"""


							super().tearDown()
							gc.collect()

					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						Any=0				, _UpperCAmelCase   :						Optional[Any]=(4, 4, 64, 64)				, _UpperCAmelCase   :						Any=False		)						->       List[Any]:




							"""simple docstring"""


							lowercase__					      =		jnp.bfloataa if fpaa else jnp.floataa
							lowercase__					      =		jnp.array(load_hf_numpy(self.get_file_format(_UpperCAmelCase				, _UpperCAmelCase		)		)				, dtype=_UpperCAmelCase		)
							return image

					def   lowerCamelCase__					(self   :						Optional[Any]				, _UpperCAmelCase   :						str=False				, _UpperCAmelCase   :						str="CompVis/stable-diffusion-v1-4"		)						->       Dict:




							"""simple docstring"""


							lowercase__					      =		jnp.bfloataa if fpaa else jnp.floataa
							lowercase__					      =		"""bf16""" if fpaa else None

							lowercase__      ,     lowercase__					      =		FlaxUNetaDConditionModel.from_pretrained(
							    _UpperCAmelCase				, subfolder="""unet"""				, dtype=_UpperCAmelCase				, revision=_UpperCAmelCase		)
							return model, params

					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						int=0				, _UpperCAmelCase   :						Optional[int]=(4, 77, 768)				, _UpperCAmelCase   :						Any=False		)						->       Optional[Any]:




							"""simple docstring"""


							lowercase__					      =		jnp.bfloataa if fpaa else jnp.floataa
							lowercase__					      =		jnp.array(load_hf_numpy(self.get_file_format(_UpperCAmelCase				, _UpperCAmelCase		)		)				, dtype=_UpperCAmelCase		)
							return hidden_states

					@parameterized.expand(
					    [
					        # fmt: off
					        [83, 4, [-0.2_323, -0.1_304, 0.0_813, -0.3_093, -0.0_919, -0.1_571, -0.1_125, -0.5_806]],
					        [17, 0.55, [-0.0_831, -0.2_443, 0.0_901, -0.0_919, 0.3_396, 0.0_103, -0.3_743, 0.0_701]],
					        [8, 0.89, [-0.4_863, 0.0_859, 0.0_875, -0.1_658, 0.9_199, -0.0_114, 0.4_839, 0.4_639]],
					        [3, 1000, [-0.5_649, 0.2_402, -0.5_518, 0.1_248, 1.1_328, -0.2_443, -0.0_325, -1.0_078]],
					        # fmt: on
					    ]		)
					def   lowerCamelCase__					(self   :						int				, _UpperCAmelCase   :						str				, _UpperCAmelCase   :						Union[str, Any]				, _UpperCAmelCase   :						int		)						->       str:




							"""simple docstring"""


							lowercase__      ,     lowercase__					      =		self.get_unet_model(model_id="""CompVis/stable-diffusion-v1-4"""				, fpaa=_UpperCAmelCase		)
							lowercase__					      =		self.get_latents(_UpperCAmelCase				, fpaa=_UpperCAmelCase		)
							lowercase__					      =		self.get_encoder_hidden_states(_UpperCAmelCase				, fpaa=_UpperCAmelCase		)

							lowercase__					      =		model.apply(
							    {"""params""": params}				, _UpperCAmelCase				, jnp.array(_UpperCAmelCase				, dtype=jnp.intaa		)				, encoder_hidden_states=_UpperCAmelCase				, ).sample

							assert sample.shape == latents.shape

							lowercase__					      =		jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())		)				, dtype=jnp.floataa		)
							lowercase__					      =		jnp.array(_UpperCAmelCase				, dtype=jnp.floataa		)

							# Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
							assert jnp.allclose(_UpperCAmelCase				, _UpperCAmelCase				, atol=1E-2		)


					@parameterized.expand(
					    [
					        # fmt: off
					        [83, 4, [0.1_514, 0.0_807, 0.1_624, 0.1_016, -0.1_896, 0.0_263, 0.0_677, 0.2_310]],
					        [17, 0.55, [0.1_164, -0.0_216, 0.0_170, 0.1_589, -0.3_120, 0.1_005, -0.0_581, -0.1_458]],
					        [8, 0.89, [-0.1_758, -0.0_169, 0.1_004, -0.1_411, 0.1_312, 0.1_103, -0.1_996, 0.2_139]],
					        [3, 1000, [0.1_214, 0.0_352, -0.0_731, -0.1_562, -0.0_994, -0.0_906, -0.2_340, -0.0_539]],
					        # fmt: on
					    ]		)
					def   lowerCamelCase__					(self   :						Optional[Any]				, _UpperCAmelCase   :						Tuple				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						List[str]		)						->       str:




							"""simple docstring"""


							lowercase__      ,     lowercase__					      =		self.get_unet_model(model_id="""stabilityai/stable-diffusion-2"""				, fpaa=_UpperCAmelCase		)
							lowercase__					      =		self.get_latents(_UpperCAmelCase				, shape=(4, 4, 96, 96)				, fpaa=_UpperCAmelCase		)
							lowercase__					      =		self.get_encoder_hidden_states(_UpperCAmelCase				, shape=(4, 77, 1024)				, fpaa=_UpperCAmelCase		)

							lowercase__					      =		model.apply(
							    {"""params""": params}				, _UpperCAmelCase				, jnp.array(_UpperCAmelCase				, dtype=jnp.intaa		)				, encoder_hidden_states=_UpperCAmelCase				, ).sample

							assert sample.shape == latents.shape

							lowercase__					      =		jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())		)				, dtype=jnp.floataa		)
							lowercase__					      =		jnp.array(_UpperCAmelCase				, dtype=jnp.floataa		)

							# Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, on the same hardware
							assert jnp.allclose(_UpperCAmelCase				, _UpperCAmelCase				, atol=1E-2		)

