





'''simple docstring'''


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


lowerCAmelCase :					int 							=logging.get_logger(__name__)

# General docstring
lowerCAmelCase :					Any 							='''RegNetConfig'''

# Base docstring
lowerCAmelCase :					Optional[Any] 							='''facebook/regnet-y-040'''
lowerCAmelCase :					Optional[int] 							=[1, 1_088, 7, 7]

# Image classification docstring
lowerCAmelCase :					Dict 							='''facebook/regnet-y-040'''
lowerCAmelCase :					List[Any] 							='''tabby, tabby cat'''

lowerCAmelCase :					Optional[int] 							=[
    '''facebook/regnet-y-040''',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]
class 				a_       (    tf.keras.layers.Layer							):





     def __init__(			self		: Optional[int]     ,				lowercase		: int     ,				lowercase		: int = 3     ,				lowercase		: int = 1     ,				lowercase		: int = 1     ,				lowercase		: Optional[str] = "relu"     ,				**lowercase		: List[str]     ,				):





        """simple docstring"""




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







     def 							lowercase__				(			self		: Any     ,				lowercase		: Union[str, Any]						):





        """simple docstring"""




        lowercase_      :Tuple         =			self.convolution(self.padding(lowercase						)						)
        lowercase_      :Any         =			self.normalization(lowercase						)
        lowercase_      :Dict         =			self.activation(lowercase						)
        return hidden_state







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





     def __init__(			self		: Dict     ,				lowercase		: RegNetConfig     ,				**lowercase		: int						):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :List[str]         =			config.num_channels
        lowercase_      :str         =			TFRegNetConvLayer(
            out_channels=config.embedding_size     ,				kernel_size=3     ,				stride=2     ,				activation=config.hidden_act     ,				name="embedder"     ,				)







     def 							lowercase__				(			self		: Dict     ,				lowercase		: int						):





        """simple docstring"""




        lowercase_      :List[Any]         =			shape_list(lowercase						)[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_      :Dict         =			tf.transpose(lowercase     ,				perm=(0, 2, 3, 1)						)
        lowercase_      :Optional[Any]         =			self.embedder(lowercase						)
        return hidden_state







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





     def __init__(			self		: str     ,				lowercase		: int     ,				lowercase		: int = 2     ,				**lowercase		: List[str]						):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :Union[str, Any]         =			tf.keras.layers.ConvaD(
            filters=lowercase     ,				kernel_size=1     ,				strides=lowercase     ,				use_bias=lowercase     ,				name="convolution"						)
        lowercase_      :Tuple         =			tf.keras.layers.BatchNormalization(epsilon=1e-5     ,				momentum=0.9     ,				name="normalization"						)







     def 							lowercase__				(			self		: str     ,				lowercase		: tf.Tensor     ,				lowercase		: bool = False						):





        """simple docstring"""




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







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





     def __init__(			self		: Dict     ,				lowercase		: int     ,				lowercase		: int     ,				**lowercase		: Union[str, Any]						):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :str         =			tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase     ,				name="pooler"						)
        lowercase_      :Optional[Any]         =			[
            tf.keras.layers.ConvaD(filters=lowercase     ,				kernel_size=1     ,				activation="relu"     ,				name="attention.0"						),
            tf.keras.layers.ConvaD(filters=lowercase     ,				kernel_size=1     ,				activation="sigmoid"     ,				name="attention.2"						),
        ]







     def 							lowercase__				(			self		: Dict     ,				lowercase		: Dict						):





        """simple docstring"""




        lowercase_      :str         =			self.pooler(lowercase						)
        for layer_module in self.attention:
           lowercase_      :Tuple         =			layer_module(lowercase						)
        lowercase_      :str         =			hidden_state * pooled
        return hidden_state







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





     def __init__(			self		: str     ,				lowercase		: RegNetConfig     ,				lowercase		: int     ,				lowercase		: int     ,				lowercase		: int = 1     ,				**lowercase		: Any						):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :List[str]         =			in_channels != out_channels or stride != 1
        lowercase_      :int         =			max(1     ,				out_channels // config.groups_width						)
        lowercase_      :List[str]         =			(
            TFRegNetShortCut(lowercase     ,				stride=lowercase     ,				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_      :str         =			[
            TFRegNetConvLayer(lowercase     ,				kernel_size=1     ,				activation=config.hidden_act     ,				name="layer.0"						),
            TFRegNetConvLayer(
                lowercase     ,				stride=lowercase     ,				groups=lowercase     ,				activation=config.hidden_act     ,				name="layer.1"						),
            TFRegNetConvLayer(lowercase     ,				kernel_size=1     ,				activation=lowercase     ,				name="layer.2"						),
        ]
        lowercase_      :Any         =			ACTaFN[config.hidden_act]







     def 							lowercase__				(			self		: Any     ,				lowercase		: str						):





        """simple docstring"""




        lowercase_      :int         =			hidden_state
        for layer_module in self.layers:
           lowercase_      :Optional[Any]         =			layer_module(lowercase						)
        lowercase_      :Dict         =			self.shortcut(lowercase						)
        hidden_state += residual
        lowercase_      :Tuple         =			self.activation(lowercase						)
        return hidden_state







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





     def __init__(			self		: Dict     ,				lowercase		: RegNetConfig     ,				lowercase		: int     ,				lowercase		: int     ,				lowercase		: int = 1     ,				**lowercase		: str						):





        """simple docstring"""




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







     def 							lowercase__				(			self		: int     ,				lowercase		: List[str]						):





        """simple docstring"""




        lowercase_      :Optional[Any]         =			hidden_state
        for layer_module in self.layers:
           lowercase_      :Tuple         =			layer_module(lowercase						)
        lowercase_      :str         =			self.shortcut(lowercase						)
        hidden_state += residual
        lowercase_      :Tuple         =			self.activation(lowercase						)
        return hidden_state







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





     def __init__(			self		: Any     ,				lowercase		: RegNetConfig     ,				lowercase		: int     ,				lowercase		: int     ,				lowercase		: int = 2     ,				lowercase		: int = 2     ,				**lowercase		: Optional[Any]						):





        """simple docstring"""




        super().__init__(**lowercase						)

        lowercase_      :Union[str, Any]         =			TFRegNetXLayer if config.layer_type == "x" else TFRegNetYLayer
        lowercase_      :Tuple         =			[
            # downsampling is done in the first layer with stride of 2
            layer(lowercase     ,				lowercase     ,				lowercase     ,				stride=lowercase     ,				name="layers.0"						),
            *[layer(lowercase     ,				lowercase     ,				lowercase     ,				name=F'layers.{i+1}'						) for i in range(depth - 1						)],
        ]







     def 							lowercase__				(			self		: Union[str, Any]     ,				lowercase		: Dict						):





        """simple docstring"""




        for layer_module in self.layers:
           lowercase_      :Any         =			layer_module(lowercase						)
        return hidden_state







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





     def __init__(			self		: Dict     ,				lowercase		: RegNetConfig     ,				**lowercase		: Optional[Any]						):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :Tuple         =			[]
        # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
        self.stages.append(
            TFRegNetStage(
                lowercase     ,				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_      :Optional[int]         =			zip(config.hidden_sizes     ,				config.hidden_sizes[1:]						)
        for i, ((in_channels, out_channels), depth) in enumerate(zip(lowercase     ,				config.depths[1:]						)						):
           self.stages.append(TFRegNetStage(lowercase     ,				lowercase     ,				lowercase     ,				depth=lowercase     ,				name=F'stages.{i+1}'						)						)







     def 							lowercase__				(			self		: List[Any]     ,				lowercase		: tf.Tensor     ,				lowercase		: bool = False     ,				lowercase		: bool = True						):





        """simple docstring"""




        lowercase_      :List[Any]         =			() if output_hidden_states else None

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

           lowercase_      :List[Any]         =			stage_module(lowercase						)

        if output_hidden_states:
           lowercase_      :Union[str, Any]         =			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=lowercase     ,				hidden_states=lowercase						)







@keras_serializable
class 				a_       (    tf.keras.layers.Layer							):
     __A								=      RegNetConfig





     def __init__(			self		: str     ,				lowercase		: Optional[int]     ,				**lowercase		: int						):





        """simple docstring"""




        super().__init__(**lowercase						)
        lowercase_      :Union[str, Any]         =			config
        lowercase_      :List[Any]         =			TFRegNetEmbeddings(lowercase     ,				name="embedder"						)
        lowercase_      :Optional[Any]         =			TFRegNetEncoder(lowercase     ,				name="encoder"						)
        lowercase_      :Dict         =			tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase     ,				name="pooler"						)







     @unpack_inputs
     def 							lowercase__				(			self		: int     ,				lowercase		: tf.Tensor     ,				lowercase		: Optional[bool] = None     ,				lowercase		: Optional[bool] = None     ,				lowercase		: bool = False     ,				):





        """simple docstring"""




        lowercase_      :List[Any]         =			(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowercase_      :Dict         =			return_dict if return_dict is not None else self.config.use_return_dict

        lowercase_      :str         =			self.embedder(lowercase     ,				training=lowercase						)

        lowercase_      :Tuple         =			self.encoder(
            lowercase     ,				output_hidden_states=lowercase     ,				return_dict=lowercase     ,				training=lowercase						)

        lowercase_      :List[str]         =			encoder_outputs[0]
        lowercase_      :Optional[Any]         =			self.pooler(lowercase						)

        # Change to NCHW output format have uniformity in the modules
        lowercase_      :List[str]         =			tf.transpose(lowercase     ,				perm=(0, 3, 1, 2)						)
        lowercase_      :Optional[int]         =			tf.transpose(lowercase     ,				perm=(0, 3, 1, 2)						)

        # Change the other hidden state outputs to NCHW as well
        if output_hidden_states:
           lowercase_      :Tuple         =			tuple([tf.transpose(lowercase     ,				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=lowercase     ,				pooler_output=lowercase     ,				hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states     ,				)







class 				a_       (    _lowerCAmelCase							):
     __A								=      RegNetConfig
     __A								=      "regnet"
     __A								=      "pixel_values"





     @property
     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




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







lowerCAmelCase :					str 							=r'''
    Parameters:
    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.
        config ([`RegNetConfig`]): 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.
'''

lowerCAmelCase :					List[Any] 							=r'''
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConveNextImageProcessor.__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.
'''
@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top."     ,						_lowerCAmelCase     ,						)
class 				a_       (    _lowerCAmelCase							):





     def __init__(			self		: Any     ,				lowercase		: RegNetConfig     ,				*lowercase		: Optional[int]     ,				**lowercase		: Any						):





        """simple docstring"""




        super().__init__(lowercase     ,				*lowercase     ,				**lowercase						)
        lowercase_      :Optional[Any]         =			TFRegNetMainLayer(lowercase     ,				name="regnet"						)







     @unpack_inputs
     @add_start_docstrings_to_model_forward(lowercase						)
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC     ,				output_type=lowercase     ,				config_class=_CONFIG_FOR_DOC     ,				modality="vision"     ,				expected_output=_EXPECTED_OUTPUT_SHAPE     ,				)
     def 							lowercase__				(			self		: Union[str, Any]     ,				lowercase		: tf.Tensor     ,				lowercase		: Optional[bool] = None     ,				lowercase		: Optional[bool] = None     ,				lowercase		: Union[str, Any]=False     ,				):





        """simple docstring"""




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

        lowercase_      :Dict         =			self.regnet(
            pixel_values=lowercase     ,				output_hidden_states=lowercase     ,				return_dict=lowercase     ,				training=lowercase     ,				)
        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(
    "\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "     ,						_lowerCAmelCase     ,						)
class 				a_       (    _lowerCAmelCase     ,						_lowerCAmelCase							):





     def __init__(			self		: Optional[Any]     ,				lowercase		: RegNetConfig     ,				*lowercase		: List[Any]     ,				**lowercase		: Optional[Any]						):





        """simple docstring"""




        super().__init__(lowercase     ,				*lowercase     ,				**lowercase						)
        lowercase_      :List[Any]         =			config.num_labels
        lowercase_      :Any         =			TFRegNetMainLayer(lowercase     ,				name="regnet"						)
        # classification head
        lowercase_      :Optional[Any]         =			[
            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(lowercase						)
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT     ,				output_type=lowercase     ,				config_class=_CONFIG_FOR_DOC     ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT     ,				)
     def 							lowercase__				(			self		: Optional[int]     ,				lowercase		: tf.Tensor = None     ,				lowercase		: tf.Tensor = None     ,				lowercase		: bool = None     ,				lowercase		: bool = None     ,				lowercase		: List[str]=False     ,				):





        """simple docstring"""




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

        lowercase_      :Any         =			self.regnet(
            lowercase     ,				output_hidden_states=lowercase     ,				return_dict=lowercase     ,				training=lowercase						)

        lowercase_      :str         =			outputs.pooler_output if return_dict else outputs[1]

        lowercase_      :Union[str, Any]         =			self.classifier[0](lowercase						)
        lowercase_      :Optional[Any]         =			self.classifier[1](lowercase						)

        lowercase_      :Union[str, Any]         =			None if labels is None else self.hf_compute_loss(labels=lowercase     ,				logits=lowercase						)

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

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









'''simple docstring'''


import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest
class 				a_       (    _lowerCAmelCase							):
     __A								=      (UniPCMultistepScheduler,)
     __A								=      (("num_inference_steps", 25),)





     def 							lowercase__				(			self		: Tuple     ,				**lowercase		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :Any         =			{
            "num_train_timesteps": 1_000,
            "beta_start": 0.00_01,
            "beta_end": 0.02,
            "beta_schedule": "linear",
            "solver_order": 2,
            "solver_type": "bh2",
        }

        config.update(**lowercase						)
        return config





     def 							lowercase__				(			self		: Optional[Any]     ,				lowercase		: str=0     ,				**lowercase		: Union[str, Any]						):





        """simple docstring"""




        lowercase_      :Tuple         =			dict(self.forward_default_kwargs						)
        lowercase_      :Tuple         =			kwargs.pop("num_inference_steps"     ,				lowercase						)
        lowercase_      :Union[str, Any]         =			self.dummy_sample
        lowercase_      :Union[str, Any]         =			0.1 * sample
        lowercase_      :Any         =			[residual + 0.2, residual + 0.15, residual + 0.10]

        for scheduler_class in self.scheduler_classes:
           lowercase_      :Tuple         =			self.get_scheduler_config(**lowercase						)
           lowercase_      :Tuple         =			scheduler_class(**lowercase						)
           scheduler.set_timesteps(lowercase						)
           # copy over dummy past residuals
           lowercase_      :str         =			dummy_past_residuals[: scheduler.config.solver_order]

           with tempfile.TemporaryDirectory() as tmpdirname:
              scheduler.save_config(lowercase						)
              lowercase_      :str         =			scheduler_class.from_pretrained(lowercase						)
              new_scheduler.set_timesteps(lowercase						)
              # copy over dummy past residuals
              lowercase_      :List[Any]         =			dummy_past_residuals[: new_scheduler.config.solver_order]

           lowercase_      ,							lowercase_      :Tuple         =			sample, sample
           for t in range(lowercase     ,				time_step + scheduler.config.solver_order + 1						):
              lowercase_      :Optional[int]         =			scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample
              lowercase_      :Union[str, Any]         =			new_scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample

              assert torch.sum(torch.abs(output - new_output						)						) < 1e-5, "Scheduler outputs are not identical"





     def 							lowercase__				(			self		: str     ,				lowercase		: Any=0     ,				**lowercase		: str						):





        """simple docstring"""




        lowercase_      :Tuple         =			dict(self.forward_default_kwargs						)
        lowercase_      :Dict         =			kwargs.pop("num_inference_steps"     ,				lowercase						)
        lowercase_      :str         =			self.dummy_sample
        lowercase_      :Optional[int]         =			0.1 * sample
        lowercase_      :Optional[int]         =			[residual + 0.2, residual + 0.15, residual + 0.10]

        for scheduler_class in self.scheduler_classes:
           lowercase_      :Any         =			self.get_scheduler_config()
           lowercase_      :Tuple         =			scheduler_class(**lowercase						)
           scheduler.set_timesteps(lowercase						)

           # copy over dummy past residuals (must be after setting timesteps)
           lowercase_      :str         =			dummy_past_residuals[: scheduler.config.solver_order]

           with tempfile.TemporaryDirectory() as tmpdirname:
              scheduler.save_config(lowercase						)
              lowercase_      :int         =			scheduler_class.from_pretrained(lowercase						)
              # copy over dummy past residuals
              new_scheduler.set_timesteps(lowercase						)

              # copy over dummy past residual (must be after setting timesteps)
              lowercase_      :Optional[Any]         =			dummy_past_residuals[: new_scheduler.config.solver_order]

           lowercase_      :Dict         =			scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample
           lowercase_      :str         =			new_scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample

           assert torch.sum(torch.abs(output - new_output						)						) < 1e-5, "Scheduler outputs are not identical"





     def 							lowercase__				(			self		: List[str]     ,				lowercase		: Optional[int]=None     ,				**lowercase		: int						):





        """simple docstring"""




        if scheduler is None:
           lowercase_      :int         =			self.scheduler_classes[0]
           lowercase_      :Optional[Any]         =			self.get_scheduler_config(**lowercase						)
           lowercase_      :Optional[Any]         =			scheduler_class(**lowercase						)

        lowercase_      :str         =			self.scheduler_classes[0]
        lowercase_      :Optional[int]         =			self.get_scheduler_config(**lowercase						)
        lowercase_      :List[Any]         =			scheduler_class(**lowercase						)

        lowercase_      :str         =			10
        lowercase_      :List[str]         =			self.dummy_model()
        lowercase_      :Optional[int]         =			self.dummy_sample_deter
        scheduler.set_timesteps(lowercase						)

        for i, t in enumerate(scheduler.timesteps						):
           lowercase_      :int         =			model(lowercase     ,				lowercase						)
           lowercase_      :Dict         =			scheduler.step(lowercase     ,				lowercase     ,				lowercase						).prev_sample

        return sample





     def 							lowercase__				(			self		: int						):





        """simple docstring"""




        lowercase_      :str         =			dict(self.forward_default_kwargs						)

        lowercase_      :Union[str, Any]         =			kwargs.pop("num_inference_steps"     ,				lowercase						)

        for scheduler_class in self.scheduler_classes:
           lowercase_      :Optional[Any]         =			self.get_scheduler_config()
           lowercase_      :Optional[int]         =			scheduler_class(**lowercase						)

           lowercase_      :Union[str, Any]         =			self.dummy_sample
           lowercase_      :str         =			0.1 * sample

           if num_inference_steps is not None and hasattr(lowercase     ,				"set_timesteps"						):
              scheduler.set_timesteps(lowercase						)
           elif num_inference_steps is not None and not hasattr(lowercase     ,				"set_timesteps"						):
              lowercase_      :str         =			num_inference_steps

           # copy over dummy past residuals (must be done after set_timesteps)
           lowercase_      :Dict         =			[residual + 0.2, residual + 0.15, residual + 0.10]
           lowercase_      :Any         =			dummy_past_residuals[: scheduler.config.solver_order]

           lowercase_      :Optional[int]         =			scheduler.timesteps[5]
           lowercase_      :List[Any]         =			scheduler.timesteps[6]

           lowercase_      :Optional[Any]         =			scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample
           lowercase_      :List[str]         =			scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample

           self.assertEqual(output_a.shape     ,				sample.shape						)
           self.assertEqual(output_a.shape     ,				output_a.shape						)





     def 							lowercase__				(			self		: List[str]						):





        """simple docstring"""




        lowercase_      :List[Any]         =			UniPCMultistepScheduler(**self.get_scheduler_config()						)
        lowercase_      :Dict         =			self.full_loop(scheduler=lowercase						)
        lowercase_      :List[Any]         =			torch.mean(torch.abs(lowercase						)						)

        assert abs(result_mean.item() - 0.24_64						) < 1e-3

        lowercase_      :List[str]         =			DPMSolverSinglestepScheduler.from_config(scheduler.config						)
        lowercase_      :Tuple         =			DEISMultistepScheduler.from_config(scheduler.config						)
        lowercase_      :Optional[int]         =			DPMSolverMultistepScheduler.from_config(scheduler.config						)
        lowercase_      :Tuple         =			UniPCMultistepScheduler.from_config(scheduler.config						)

        lowercase_      :Tuple         =			self.full_loop(scheduler=lowercase						)
        lowercase_      :Tuple         =			torch.mean(torch.abs(lowercase						)						)

        assert abs(result_mean.item() - 0.24_64						) < 1e-3





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        for timesteps in [25, 50, 100, 999, 1_000]:
           self.check_over_configs(num_train_timesteps=lowercase						)





     def 							lowercase__				(			self		: int						):





        """simple docstring"""




        self.check_over_configs(thresholding=lowercase						)
        for order in [1, 2, 3]:
           for solver_type in ["bh1", "bh2"]:
              for threshold in [0.5, 1.0, 2.0]:
                 for prediction_type in ["epsilon", "sample"]:
                    self.check_over_configs(
                        thresholding=lowercase     ,				prediction_type=lowercase     ,				sample_max_value=lowercase     ,				solver_order=lowercase     ,				solver_type=lowercase     ,				)





     def 							lowercase__				(			self		: List[str]						):





        """simple docstring"""




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





     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        for solver_type in ["bh1", "bh2"]:
           for order in [1, 2, 3]:
              for prediction_type in ["epsilon", "sample"]:
                 self.check_over_configs(
                     solver_order=lowercase     ,				solver_type=lowercase     ,				prediction_type=lowercase     ,				)
                 lowercase_      :Tuple         =			self.full_loop(
                     solver_order=lowercase     ,				solver_type=lowercase     ,				prediction_type=lowercase     ,				)
                 assert not torch.isnan(lowercase						).any(), "Samples have nan numbers"





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        self.check_over_configs(lower_order_final=lowercase						)
        self.check_over_configs(lower_order_final=lowercase						)





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1_000]:
           self.check_over_forward(num_inference_steps=lowercase     ,				time_step=0						)





     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        lowercase_      :List[str]         =			self.full_loop()
        lowercase_      :List[str]         =			torch.mean(torch.abs(lowercase						)						)

        assert abs(result_mean.item() - 0.24_64						) < 1e-3





     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        lowercase_      :Dict         =			self.full_loop(prediction_type="v_prediction"						)
        lowercase_      :Any         =			torch.mean(torch.abs(lowercase						)						)

        assert abs(result_mean.item() - 0.10_14						) < 1e-3





     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			self.scheduler_classes[0]
        lowercase_      :str         =			self.get_scheduler_config(thresholding=lowercase     ,				dynamic_thresholding_ratio=0						)
        lowercase_      :Tuple         =			scheduler_class(**lowercase						)

        lowercase_      :Optional[int]         =			10
        lowercase_      :Optional[int]         =			self.dummy_model()
        lowercase_      :Any         =			self.dummy_sample_deter.half()
        scheduler.set_timesteps(lowercase						)

        for i, t in enumerate(scheduler.timesteps						):
           lowercase_      :Tuple         =			model(lowercase     ,				lowercase						)
           lowercase_      :Any         =			scheduler.step(lowercase     ,				lowercase     ,				lowercase						).prev_sample

        assert sample.dtype == torch.floataa







     def 							lowercase__				(			self		: List[Any]     ,				**lowercase		: Union[str, Any]						):





        """simple docstring"""




        for scheduler_class in self.scheduler_classes:
           lowercase_      :Optional[int]         =			self.get_scheduler_config(**lowercase						)
           lowercase_      :Dict         =			scheduler_class(**lowercase						)

           scheduler.set_timesteps(scheduler.config.num_train_timesteps						)
           assert len(scheduler.timesteps.unique()						) == scheduler.num_inference_steps


