




from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


lowercase_		=							logging.get_logger(__name__)

# General docstring
lowercase_		=							"""ResNetConfig"""

# Base docstring
lowercase_		=							"""microsoft/resnet-50"""
lowercase_		=							[1, 2048, 7, 7]

# Image classification docstring
lowercase_		=							"""microsoft/resnet-50"""
lowercase_		=							"""tiger cat"""

lowercase_		=							[
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     lowerCamelCase = 3				,     lowerCamelCase = 1				,     lowerCamelCase = "relu"				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)
						__magic_name__					:							Any		   =							kernel_size // 2
						__magic_name__					:							int		   =							tf.keras.layers.ConvaD(
						    lowerCamelCase__				,     kernel_size=lowerCamelCase__				,     strides=lowerCamelCase__				,     padding='''valid'''				,     use_bias=lowerCamelCase__				,     name='''convolution'''						)
						# Use same default momentum and epsilon as PyTorch equivalent
						__magic_name__					:							Tuple		   =							tf.keras.layers.BatchNormalization(epsilon=1e-5				,     momentum=0.9				,     name='''normalization'''						)
						__magic_name__					:							Any		   =							ACTaFN[activation] if activation is not None else tf.keras.layers.Activation('''linear'''						)






	def  lowercase      (	self				,     lowerCamelCase						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__					:							int		   =							(self.pad_value, self.pad_value)
						__magic_name__					:							Optional[Any]		   =							tf.pad(lowerCamelCase__				,     [(0, 0), height_pad, width_pad, (0, 0)]						)
						__magic_name__					:							Optional[Any]		   =							self.conv(lowerCamelCase__						)
						return hidden_state






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__					:							Tuple		   =							self.convolution(lowerCamelCase__						)
						__magic_name__					:							Any		   =							self.normalization(lowerCamelCase__				,     training=lowerCamelCase__						)
						__magic_name__					:							int		   =							self.activation(lowerCamelCase__						)
						return hidden_state





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)
						__magic_name__					:							Optional[Any]		   =							TFResNetConvLayer(
						    config.embedding_size				,     kernel_size=7				,     stride=2				,     activation=config.hidden_act				,     name='''embedder'''				,     )
						__magic_name__					:							int		   =							tf.keras.layers.MaxPoolaD(pool_size=3				,     strides=2				,     padding='''valid'''				,     name='''pooler'''						)
						__magic_name__					:							Optional[int]		   =							config.num_channels






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__   , __magic_name__   , __magic_name__   , __magic_name__					:							Tuple		   =							shape_list(lowerCamelCase__						)
						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.'''						)
						__magic_name__					:							Optional[int]		   =							pixel_values
						__magic_name__					:							List[Any]		   =							self.embedder(lowerCamelCase__						)
						__magic_name__					:							Dict		   =							tf.pad(lowerCamelCase__				,     [[0, 0], [1, 1], [1, 1], [0, 0]]						)
						__magic_name__					:							List[str]		   =							self.pooler(lowerCamelCase__						)
						return hidden_state





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     lowerCamelCase = 2				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)
						__magic_name__					:							str		   =							tf.keras.layers.ConvaD(
						    lowerCamelCase__				,     kernel_size=1				,     strides=lowerCamelCase__				,     use_bias=lowerCamelCase__				,     name='''convolution'''						)
						# Use same default momentum and epsilon as PyTorch equivalent
						__magic_name__					:							Tuple		   =							tf.keras.layers.BatchNormalization(epsilon=1e-5				,     momentum=0.9				,     name='''normalization'''						)






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__					:							Optional[Any]		   =							x
						__magic_name__					:							Dict		   =							self.convolution(lowerCamelCase__						)
						__magic_name__					:							str		   =							self.normalization(lowerCamelCase__				,     training=lowerCamelCase__						)
						return hidden_state





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase = 1				,     lowerCamelCase = "relu"				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)
						__magic_name__					:							Optional[Any]		   =							in_channels != out_channels or stride != 1
						__magic_name__					:							str		   =							TFResNetConvLayer(lowerCamelCase__				,     stride=lowerCamelCase__				,     name='''layer.0'''						)
						__magic_name__					:							Dict		   =							TFResNetConvLayer(lowerCamelCase__				,     activation=lowerCamelCase__				,     name='''layer.1'''						)
						__magic_name__					:							Optional[Any]		   =							(
						    TFResNetShortCut(lowerCamelCase__				,     stride=lowerCamelCase__				,     name='''shortcut'''						)
						    if should_apply_shortcut
						    else tf.keras.layers.Activation('''linear'''				,     name='''shortcut'''						)
						)
						__magic_name__					:							List[Any]		   =							ACTaFN[activation]






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__					:							Union[str, Any]		   =							hidden_state
						__magic_name__					:							Optional[Any]		   =							self.conva(lowerCamelCase__				,     training=lowerCamelCase__						)
						__magic_name__					:							Tuple		   =							self.conva(lowerCamelCase__				,     training=lowerCamelCase__						)
						__magic_name__					:							Optional[int]		   =							self.shortcut(lowerCamelCase__				,     training=lowerCamelCase__						)
						hidden_state += residual
						__magic_name__					:							str		   =							self.activation(lowerCamelCase__						)
						return hidden_state





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase = 1				,     lowerCamelCase = "relu"				,     lowerCamelCase = 4				,     **lowerCamelCase				,     )      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)
						__magic_name__					:							Union[str, Any]		   =							in_channels != out_channels or stride != 1
						__magic_name__					:							str		   =							out_channels // reduction
						__magic_name__					:							List[Any]		   =							TFResNetConvLayer(lowerCamelCase__				,     kernel_size=1				,     name='''layer.0'''						)
						__magic_name__					:							Union[str, Any]		   =							TFResNetConvLayer(lowerCamelCase__				,     stride=lowerCamelCase__				,     name='''layer.1'''						)
						__magic_name__					:							Optional[int]		   =							TFResNetConvLayer(lowerCamelCase__				,     kernel_size=1				,     activation=lowerCamelCase__				,     name='''layer.2'''						)
						__magic_name__					:							str		   =							(
						    TFResNetShortCut(lowerCamelCase__				,     stride=lowerCamelCase__				,     name='''shortcut'''						)
						    if should_apply_shortcut
						    else tf.keras.layers.Activation('''linear'''				,     name='''shortcut'''						)
						)
						__magic_name__					:							Union[str, Any]		   =							ACTaFN[activation]






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__					:							List[str]		   =							hidden_state
						__magic_name__					:							Any		   =							self.conva(lowerCamelCase__				,     training=lowerCamelCase__						)
						__magic_name__					:							Union[str, Any]		   =							self.conva(lowerCamelCase__				,     training=lowerCamelCase__						)
						__magic_name__					:							Optional[int]		   =							self.conva(lowerCamelCase__				,     training=lowerCamelCase__						)
						__magic_name__					:							Any		   =							self.shortcut(lowerCamelCase__				,     training=lowerCamelCase__						)
						hidden_state += residual
						__magic_name__					:							Union[str, Any]		   =							self.activation(lowerCamelCase__						)
						return hidden_state





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase = 2				,     lowerCamelCase = 2				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)

						__magic_name__					:							List[Any]		   =							TFResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else TFResNetBasicLayer

						__magic_name__					:							int		   =							[layer(lowerCamelCase__				,     lowerCamelCase__				,     stride=lowerCamelCase__				,     activation=config.hidden_act				,     name='''layers.0'''						)]
						layers += [
						    layer(lowerCamelCase__				,     lowerCamelCase__				,     activation=config.hidden_act				,     name=F'''layers.{i + 1}'''						)
						    for i in range(depth - 1						)
						]
						__magic_name__					:							Union[str, Any]		   =							layers






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False						)      ->   tf.Tensor:




						"""simple docstring"""






						for layer in self.stage_layers:
											__magic_name__					:							Tuple		   =							layer(lowerCamelCase__				,     training=lowerCamelCase__						)
						return hidden_state





class 		A__				(   tf.keras.layers.Layer ):






	def __init__(	self				,     lowerCamelCase				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






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






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = False				,     lowerCamelCase = True				,     lowerCamelCase = False				,     )      ->   TFBaseModelOutputWithNoAttention:




						"""simple docstring"""






						__magic_name__					:							List[str]		   =							() if output_hidden_states else None

						for stage_module in self.stages:
											if output_hidden_states:
																__magic_name__					:							Union[str, Any]		   =							hidden_states + (hidden_state,)

											__magic_name__					:							Dict		   =							stage_module(lowerCamelCase__				,     training=lowerCamelCase__						)

						if output_hidden_states:
											__magic_name__					:							Tuple		   =							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=lowerCamelCase__				,     hidden_states=lowerCamelCase__						)





class 		A__				(   __A ):
	lowerCamelCase__		: Dict				    =ResNetConfig
	lowerCamelCase__		: str				    ="""resnet"""
	lowerCamelCase__		: List[str]				    ="""pixel_values"""






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




						"""simple docstring"""






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

lowercase_		=							r"""
    This model is a TensorFlow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""


lowercase_		=							r"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""





@keras_serializable
class 		A__				(   tf.keras.layers.Layer ):
	lowerCamelCase__		: List[str]				    =ResNetConfig






	def __init__(	self				,     lowerCamelCase				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(**lowerCamelCase__						)
						__magic_name__					:							int		   =							config
						__magic_name__					:							List[Any]		   =							TFResNetEmbeddings(lowerCamelCase__				,     name='''embedder'''						)
						__magic_name__					:							Optional[Any]		   =							TFResNetEncoder(lowerCamelCase__				,     name='''encoder'''						)
						__magic_name__					:							List[Any]		   =							tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowerCamelCase__						)






	@unpack_inputs
	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = None				,     lowerCamelCase = None				,     lowerCamelCase = False				,     )      ->   Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:




						"""simple docstring"""






						__magic_name__					:							int		   =							(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						__magic_name__					:							List[str]		   =							return_dict if return_dict is not None else self.config.use_return_dict

						# TF 2.0 image layers can't use NCHW format when running on CPU.
						# We transpose to NHWC format and then transpose back after the full forward pass.
						# (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
						__magic_name__					:							int		   =							tf.transpose(lowerCamelCase__				,     perm=[0, 2, 3, 1]						)
						__magic_name__					:							Tuple		   =							self.embedder(lowerCamelCase__				,     training=lowerCamelCase__						)

						__magic_name__					:							Dict		   =							self.encoder(
						    lowerCamelCase__				,     output_hidden_states=lowerCamelCase__				,     return_dict=lowerCamelCase__				,     training=lowerCamelCase__						)

						__magic_name__					:							int		   =							encoder_outputs[0]

						__magic_name__					:							Any		   =							self.pooler(lowerCamelCase__						)

						# Transpose all the outputs to the NCHW format
						# (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
						__magic_name__					:							Optional[Any]		   =							tf.transpose(lowerCamelCase__				,     (0, 3, 1, 2)						)
						__magic_name__					:							Optional[Any]		   =							tf.transpose(lowerCamelCase__				,     (0, 3, 1, 2)						)
						__magic_name__					:							Optional[Any]		   =							()
						for hidden_state in encoder_outputs[1:]:
											__magic_name__					:							int		   =							hidden_states + tuple(tf.transpose(lowerCamelCase__				,     (0, 3, 1, 2)						) for h in hidden_state						)

						if not return_dict:
											return (last_hidden_state, pooled_output) + hidden_states

						__magic_name__					:							List[Any]		   =							hidden_states if output_hidden_states else None

						return TFBaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=lowerCamelCase__				,     pooler_output=lowerCamelCase__				,     hidden_states=lowerCamelCase__				,     )





@add_start_docstrings(
    "The bare ResNet model outputting raw features without any specific head on top."		,     __A		,     )
class 		A__				(   __A ):






	def __init__(	self				,     lowerCamelCase				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(lowerCamelCase__				,     **lowerCamelCase__						)
						__magic_name__					:							int		   =							TFResNetMainLayer(config=lowerCamelCase__				,     name='''resnet'''						)






	@add_start_docstrings_to_model_forward(lowerCamelCase__						)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC				,     output_type=lowerCamelCase__				,     config_class=_CONFIG_FOR_DOC				,     modality='''vision'''				,     expected_output=_EXPECTED_OUTPUT_SHAPE				,     )
	@unpack_inputs
	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase = None				,     lowerCamelCase = None				,     lowerCamelCase = False				,     )      ->   Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:




						"""simple docstring"""






						__magic_name__					:							Tuple		   =							(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						__magic_name__					:							Union[str, Any]		   =							return_dict if return_dict is not None else self.config.use_return_dict

						__magic_name__					:							List[Any]		   =							self.resnet(
						    pixel_values=lowerCamelCase__				,     output_hidden_states=lowerCamelCase__				,     return_dict=lowerCamelCase__				,     training=lowerCamelCase__				,     )
						return resnet_outputs





@add_start_docstrings(
    "\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "		,     __A		,     )
class 		A__				(   __A		,     __A ):






	def __init__(	self				,     lowerCamelCase				,     **lowerCamelCase						)      ->   None:




						"""simple docstring"""






						super().__init__(lowerCamelCase__				,     **lowerCamelCase__						)
						__magic_name__					:							Any		   =							config.num_labels
						__magic_name__					:							Any		   =							TFResNetMainLayer(lowerCamelCase__				,     name='''resnet'''						)
						# classification head
						__magic_name__					:							Optional[Any]		   =							(
						    tf.keras.layers.Dense(config.num_labels				,     name='''classifier.1'''						)
						    if config.num_labels > 0
						    else tf.keras.layers.Activation('''linear'''				,     name='''classifier.1'''						)
						)






	def  lowercase      (	self				,     lowerCamelCase						)      ->   tf.Tensor:




						"""simple docstring"""






						__magic_name__					:							Optional[Any]		   =							tf.keras.layers.Flatten()(lowerCamelCase__						)
						__magic_name__					:							Union[str, Any]		   =							self.classifier_layer(lowerCamelCase__						)
						return logits






	@add_start_docstrings_to_model_forward(lowerCamelCase__						)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT				,     output_type=lowerCamelCase__				,     config_class=_CONFIG_FOR_DOC				,     expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,     )
	@unpack_inputs
	def  lowercase      (	self				,     lowerCamelCase = None				,     lowerCamelCase = None				,     lowerCamelCase = None				,     lowerCamelCase = None				,     lowerCamelCase = False				,     )      ->   Union[Tuple[tf.Tensor], TFImageClassifierOutputWithNoAttention]:




						"""simple docstring"""






						__magic_name__					:							List[str]		   =							return_dict if return_dict is not None else self.config.use_return_dict

						__magic_name__					:							List[Any]		   =							self.resnet(
						    lowerCamelCase__				,     output_hidden_states=lowerCamelCase__				,     return_dict=lowerCamelCase__				,     training=lowerCamelCase__						)

						__magic_name__					:							Dict		   =							outputs.pooler_output if return_dict else outputs[1]

						__magic_name__					:							Optional[int]		   =							self.classifier(lowerCamelCase__						)

						__magic_name__					:							Dict		   =							None if labels is None else self.hf_compute_loss(lowerCamelCase__				,     lowerCamelCase__						)

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

						return TFImageClassifierOutputWithNoAttention(loss=lowerCamelCase__				,     logits=lowerCamelCase__				,     hidden_states=outputs.hidden_states						)









'''simple docstring'''
import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest






class        SCREAMING_SNAKE_CASE(				__A     ):
       snake_case_ :     Any    =					(HeunDiscreteScheduler,)
       snake_case_ :     Any    =					10

       def     snake_case__			(    self					,		**lowerCamelCase__     )    ->					int:

           """simple docstring"""



           __lowercase							=						{
               """num_train_timesteps""": 1100,
               """beta_start""": 0.00_01,
               """beta_end""": 0.02,
               """beta_schedule""": """linear""",
           }

           config.update(**lowerCamelCase__     )
           return config

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

           """simple docstring"""



           for timesteps in [10, 50, 100, 1000]:
               self.check_over_configs(num_train_timesteps=lowerCamelCase__     )

       def     snake_case__			(    self     )    ->					Dict:

           """simple docstring"""



           for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]					,		[0.00_02, 0.0_02, 0.02]     ):
               self.check_over_configs(beta_start=lowerCamelCase__					,		beta_end=lowerCamelCase__     )

       def     snake_case__			(    self     )    ->					Any:

           """simple docstring"""



           for schedule in ["linear", "scaled_linear", "exp"]:
               self.check_over_configs(beta_schedule=lowerCamelCase__     )

       def     snake_case__			(    self     )    ->					int:

           """simple docstring"""



           for clip_sample_range in [1.0, 2.0, 3.0]:
               self.check_over_configs(clip_sample_range=lowerCamelCase__					,		clip_sample=lowerCamelCase__     )

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

           """simple docstring"""



           for prediction_type in ["epsilon", "v_prediction", "sample"]:
               self.check_over_configs(prediction_type=lowerCamelCase__     )

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase							=						self.scheduler_classes[0]
           __lowercase							=						self.get_scheduler_config()
           __lowercase							=						scheduler_class(**lowerCamelCase__     )

           scheduler.set_timesteps(self.num_inference_steps     )

           __lowercase							=						self.dummy_model()
           __lowercase							=						self.dummy_sample_deter * scheduler.init_noise_sigma
           __lowercase							=						sample.to(lowerCamelCase__     )

           for i, t in enumerate(scheduler.timesteps     ):
               __lowercase							=						scheduler.scale_model_input(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						model(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						scheduler.step(lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )
               __lowercase							=						output.prev_sample

           __lowercase							=						torch.sum(torch.abs(lowerCamelCase__     )     )
           __lowercase							=						torch.mean(torch.abs(lowerCamelCase__     )     )

           if torch_device in ["cpu", "mps"]:
               assert abs(result_sum.item() - 0.12_33     ) < 1E-2
               assert abs(result_mean.item() - 0.00_02     ) < 1E-3
           else:
               # CUDA
               assert abs(result_sum.item() - 0.12_33     ) < 1E-2
               assert abs(result_mean.item() - 0.00_02     ) < 1E-3

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase							=						self.scheduler_classes[0]
           __lowercase							=						self.get_scheduler_config(prediction_type="""v_prediction"""     )
           __lowercase							=						scheduler_class(**lowerCamelCase__     )

           scheduler.set_timesteps(self.num_inference_steps     )

           __lowercase							=						self.dummy_model()
           __lowercase							=						self.dummy_sample_deter * scheduler.init_noise_sigma
           __lowercase							=						sample.to(lowerCamelCase__     )

           for i, t in enumerate(scheduler.timesteps     ):
               __lowercase							=						scheduler.scale_model_input(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						model(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						scheduler.step(lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )
               __lowercase							=						output.prev_sample

           __lowercase							=						torch.sum(torch.abs(lowerCamelCase__     )     )
           __lowercase							=						torch.mean(torch.abs(lowerCamelCase__     )     )

           if torch_device in ["cpu", "mps"]:
               assert abs(result_sum.item() - 4.6_9_3_4E-0_7     ) < 1E-2
               assert abs(result_mean.item() - 6.1_1_1_2E-1_0     ) < 1E-3
           else:
               # CUDA
               assert abs(result_sum.item() - 4.6_9_3_4_2_8_6_5_0_1_7_0_9_7_2E-0_7     ) < 1E-2
               assert abs(result_mean.item() - 0.00_02     ) < 1E-3

       def     snake_case__			(    self     )    ->					Dict:

           """simple docstring"""



           __lowercase							=						self.scheduler_classes[0]
           __lowercase							=						self.get_scheduler_config()
           __lowercase							=						scheduler_class(**lowerCamelCase__     )

           scheduler.set_timesteps(self.num_inference_steps					,		device=lowerCamelCase__     )

           __lowercase							=						self.dummy_model()
           __lowercase							=						self.dummy_sample_deter.to(lowerCamelCase__     ) * scheduler.init_noise_sigma

           for t in scheduler.timesteps:
               __lowercase							=						scheduler.scale_model_input(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						model(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						scheduler.step(lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )
               __lowercase							=						output.prev_sample

           __lowercase							=						torch.sum(torch.abs(lowerCamelCase__     )     )
           __lowercase							=						torch.mean(torch.abs(lowerCamelCase__     )     )

           if str(lowerCamelCase__     ).startswith("""cpu"""     ):
               # The following sum varies between 148 and 156 on mps. Why?
               assert abs(result_sum.item() - 0.12_33     ) < 1E-2
               assert abs(result_mean.item() - 0.00_02     ) < 1E-3
           elif str(lowerCamelCase__     ).startswith("""mps"""     ):
               # Larger tolerance on mps
               assert abs(result_mean.item() - 0.00_02     ) < 1E-2
           else:
               # CUDA
               assert abs(result_sum.item() - 0.12_33     ) < 1E-2
               assert abs(result_mean.item() - 0.00_02     ) < 1E-3




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

           """simple docstring"""



           __lowercase							=						self.scheduler_classes[0]
           __lowercase							=						self.get_scheduler_config()
           __lowercase							=						scheduler_class(**lowerCamelCase__					,		use_karras_sigmas=lowerCamelCase__     )

           scheduler.set_timesteps(self.num_inference_steps					,		device=lowerCamelCase__     )

           __lowercase							=						self.dummy_model()
           __lowercase							=						self.dummy_sample_deter.to(lowerCamelCase__     ) * scheduler.init_noise_sigma
           __lowercase							=						sample.to(lowerCamelCase__     )

           for t in scheduler.timesteps:
               __lowercase							=						scheduler.scale_model_input(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						model(lowerCamelCase__					,		lowerCamelCase__     )

               __lowercase							=						scheduler.step(lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )
               __lowercase							=						output.prev_sample

           __lowercase							=						torch.sum(torch.abs(lowerCamelCase__     )     )
           __lowercase							=						torch.mean(torch.abs(lowerCamelCase__     )     )

           assert abs(result_sum.item() - 0.0_00_15     ) < 1E-2
           assert abs(result_mean.item() - 1.9_8_6_9_5_5_4_5_3_5_0_3_4_6_9_5E-0_7     ) < 1E-2




