




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_					    = logging.get_logger(__name__)

# General docstring
lowerCAmelCase_					    = """RegNetConfig"""

# Base docstring
lowerCAmelCase_					    = """facebook/regnet-y-040"""
lowerCAmelCase_					    = [1, 10_88, 7, 7]

# Image classification docstring
lowerCAmelCase_					    = """facebook/regnet-y-040"""
lowerCAmelCase_					    = """tabby, tabby cat"""

lowerCAmelCase_					    = [
    """facebook/regnet-y-040""",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]






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

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







      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  str		       =     self.convolution(self.padding(__UpperCAmelCase		)		)
        lowerCAmelCase__				:  List[str]		       =     self.normalization(__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     self.activation(__UpperCAmelCase		)
        return hidden_state






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

      def __init__(   self      ,				__UpperCAmelCase      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  List[str]		       =     config.num_channels
        lowerCAmelCase__				:  str		       =     TFRegNetConvLayer(
            out_channels=config.embedding_size      ,				kernel_size=3      ,				stride=2      ,				activation=config.hidden_act      ,				name='''embedder'''      ,				)







      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  int		       =     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)
        lowerCAmelCase__				:  str		       =     tf.transpose(__UpperCAmelCase      ,				perm=(0, 2, 3, 1)		)
        lowerCAmelCase__				:  Union[str, Any]		       =     self.embedder(__UpperCAmelCase		)
        return hidden_state






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

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase = 2      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     tf.keras.layers.ConvaD(
            filters=__UpperCAmelCase      ,				kernel_size=1      ,				strides=__UpperCAmelCase      ,				use_bias=__UpperCAmelCase      ,				name='''convolution'''		)
        lowerCAmelCase__				:  Tuple		       =     tf.keras.layers.BatchNormalization(epsilon=1e-5      ,				momentum=0.9      ,				name='''normalization'''		)







      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase = False		):
        return self.normalization(self.convolution(__UpperCAmelCase		)      ,				training=__UpperCAmelCase		)






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

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     tf.keras.layers.GlobalAveragePoolingaD(keepdims=__UpperCAmelCase      ,				name='''pooler'''		)
        lowerCAmelCase__				:  Optional[Any]		       =     [
            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 	__magic_name__(   self      ,				__UpperCAmelCase		):
        # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
        lowerCAmelCase__				:  Optional[Any]		       =     self.pooler(__UpperCAmelCase		)
        for layer_module in self.attention:
          lowerCAmelCase__				:  Tuple		       =     layer_module(__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[Any]		       =     hidden_state * pooled
        return hidden_state






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

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase = 1      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     in_channels != out_channels or stride != 1
        lowerCAmelCase__				:  Optional[int]		       =     max(1      ,				out_channels // config.groups_width		)
        lowerCAmelCase__				:  Dict		       =     (
            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.
        lowerCAmelCase__				:  Optional[int]		       =     [
            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'''		),
        ]
        lowerCAmelCase__				:  Optional[Any]		       =     ACTaFN[config.hidden_act]







      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  List[str]		       =     hidden_state
        for layer_module in self.layers:
          lowerCAmelCase__				:  Dict		       =     layer_module(__UpperCAmelCase		)
        lowerCAmelCase__				:  Tuple		       =     self.shortcut(__UpperCAmelCase		)
        hidden_state += residual
        lowerCAmelCase__				:  Union[str, Any]		       =     self.activation(__UpperCAmelCase		)
        return hidden_state






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

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase = 1      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  int		       =     in_channels != out_channels or stride != 1
        lowerCAmelCase__				:  Union[str, Any]		       =     max(1      ,				out_channels // config.groups_width		)
        lowerCAmelCase__				:  Optional[int]		       =     (
            TFRegNetShortCut(__UpperCAmelCase      ,				stride=__UpperCAmelCase      ,				name='''shortcut'''		)
            if should_apply_shortcut
            else tf.keras.layers.Activation('''linear'''      ,				name='''shortcut'''		)
        )
        lowerCAmelCase__				:  Optional[Any]		       =     [
            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'''		),
        ]
        lowerCAmelCase__				:  Optional[int]		       =     ACTaFN[config.hidden_act]







      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Union[str, Any]		       =     hidden_state
        for layer_module in self.layers:
          lowerCAmelCase__				:  int		       =     layer_module(__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     self.shortcut(__UpperCAmelCase		)
        hidden_state += residual
        lowerCAmelCase__				:  Any		       =     self.activation(__UpperCAmelCase		)
        return hidden_state






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

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase = 2      ,				__UpperCAmelCase = 2      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)

        lowerCAmelCase__				:  int		       =     TFRegNetXLayer if config.layer_type == '''x''' else TFRegNetYLayer
        lowerCAmelCase__				:  Optional[int]		       =     [
            # 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 	__magic_name__(   self      ,				__UpperCAmelCase		):
        for layer_module in self.layers:
          lowerCAmelCase__				:  int		       =     layer_module(__UpperCAmelCase		)
        return hidden_state






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

      def __init__(   self      ,				__UpperCAmelCase      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  int		       =     []
        # 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'''      ,				)		)
        lowerCAmelCase__				:  Union[str, Any]		       =     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 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase = False      ,				__UpperCAmelCase = True		):
        lowerCAmelCase__				:  str		       =     () if output_hidden_states else None

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

          lowerCAmelCase__				:  Dict		       =     stage_module(__UpperCAmelCase		)

        if output_hidden_states:
          lowerCAmelCase__				:  List[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=__UpperCAmelCase      ,				hidden_states=__UpperCAmelCase		)






@keras_serializable
class  _lowerCAmelCase     (  tf.keras.layers.Layer	):
      A__		     =		RegNetConfig

      def __init__(   self      ,				__UpperCAmelCase      ,				**__UpperCAmelCase		):
        super().__init__(**__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[Any]		       =     config
        lowerCAmelCase__				:  List[str]		       =     TFRegNetEmbeddings(__UpperCAmelCase      ,				name='''embedder'''		)
        lowerCAmelCase__				:  Tuple		       =     TFRegNetEncoder(__UpperCAmelCase      ,				name='''encoder'''		)
        lowerCAmelCase__				:  Optional[Any]		       =     tf.keras.layers.GlobalAveragePoolingaD(keepdims=__UpperCAmelCase      ,				name='''pooler'''		)







      @unpack_inputs
      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase = False      ,				):
        lowerCAmelCase__				:  Union[str, Any]		       =     (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowerCAmelCase__				:  str		       =     return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase__				:  int		       =     self.embedder(__UpperCAmelCase      ,				training=__UpperCAmelCase		)

        lowerCAmelCase__				:  str		       =     self.encoder(
            __UpperCAmelCase      ,				output_hidden_states=__UpperCAmelCase      ,				return_dict=__UpperCAmelCase      ,				training=__UpperCAmelCase		)

        lowerCAmelCase__				:  int		       =     encoder_outputs[0]
        lowerCAmelCase__				:  int		       =     self.pooler(__UpperCAmelCase		)

        # Change to NCHW output format have uniformity in the modules
        lowerCAmelCase__				:  List[Any]		       =     tf.transpose(__UpperCAmelCase      ,				perm=(0, 3, 1, 2)		)
        lowerCAmelCase__				:  List[Any]		       =     tf.transpose(__UpperCAmelCase      ,				perm=(0, 3, 1, 2)		)

        # Change the other hidden state outputs to NCHW as well
        if output_hidden_states:
          lowerCAmelCase__				:  Dict		       =     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  _lowerCAmelCase     (  _lowercase	):
      A__		     =		RegNetConfig
      A__		     =		'regnet'
      A__		     =		'pixel_values'

      @property
      def 	__magic_name__(   self		):
        return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)      ,				dtype=tf.floataa		)}
lowerCAmelCase_					    = 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_					    = 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.'    ,			_lowercase    ,			)
class  _lowerCAmelCase     (  _lowercase	):

      def __init__(   self      ,				__UpperCAmelCase      ,				*__UpperCAmelCase      ,				**__UpperCAmelCase		):
        super().__init__(__UpperCAmelCase      ,				*__UpperCAmelCase      ,				**__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     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 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase=False      ,				):
        lowerCAmelCase__				:  List[Any]		       =     (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowerCAmelCase__				:  List[Any]		       =     return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase__				:  Optional[Any]		       =     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(
    '\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    '    ,			_lowercase    ,			)
class  _lowerCAmelCase     (  _lowercase    ,			_lowercase	):

      def __init__(   self      ,				__UpperCAmelCase      ,				*__UpperCAmelCase      ,				**__UpperCAmelCase		):
        super().__init__(__UpperCAmelCase      ,				*__UpperCAmelCase      ,				**__UpperCAmelCase		)
        lowerCAmelCase__				:  Any		       =     config.num_labels
        lowerCAmelCase__				:  Dict		       =     TFRegNetMainLayer(__UpperCAmelCase      ,				name='''regnet'''		)
        # classification head
        lowerCAmelCase__				:  Union[str, 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(__UpperCAmelCase		)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT      ,				output_type=__UpperCAmelCase      ,				config_class=_CONFIG_FOR_DOC      ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,				)
      def 	__magic_name__(   self      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase = None      ,				__UpperCAmelCase=False      ,				):
        lowerCAmelCase__				:  Dict		       =     (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowerCAmelCase__				:  Dict		       =     return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase__				:  Dict		       =     self.regnet(
            __UpperCAmelCase      ,				output_hidden_states=__UpperCAmelCase      ,				return_dict=__UpperCAmelCase      ,				training=__UpperCAmelCase		)

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

        lowerCAmelCase__				:  List[str]		       =     self.classifier[0](__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     self.classifier[1](__UpperCAmelCase		)

        lowerCAmelCase__				:  List[Any]		       =     None if labels is None else self.hf_compute_loss(labels=__UpperCAmelCase      ,				logits=__UpperCAmelCase		)

        if not return_dict:
          lowerCAmelCase__				:  int		       =     (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 baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler






class  _lowerCAmelCase     :

      def __init__(   self      ,				__UpperCAmelCase		):
        if isinstance(__UpperCAmelCase      ,				__UpperCAmelCase		):
          # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
          # modified it, it will not be accepted here again, since `auto` values would have been overridden
          lowerCAmelCase__				:  Union[str, Any]		       =     deepcopy(__UpperCAmelCase		)
        elif os.path.exists(__UpperCAmelCase		):
          with io.open(__UpperCAmelCase      ,				'''r'''      ,				encoding='''utf-8'''		) as f:
            lowerCAmelCase__				:  int		       =     json.load(__UpperCAmelCase		)
        else:
          try:
            lowerCAmelCase__				:  Union[str, Any]		       =     baseaa.urlsafe_baadecode(__UpperCAmelCase		).decode('''utf-8'''		)
            lowerCAmelCase__				:  Optional[int]		       =     json.loads(__UpperCAmelCase		)
          except (UnicodeDecodeError, AttributeError, ValueError):
            raise ValueError(
                f"""Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}"""		)

        lowerCAmelCase__				:  Optional[Any]		       =     config

        self.set_stage_and_offload()

      def 	__magic_name__(   self		):
        # zero stage - this is done as early as possible, before model is created, to allow
        # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
        # during ``zero.Init()`` which needs to know the dtype, and some other hparams.
        lowerCAmelCase__				:  Optional[Any]		       =     self.get_value('''zero_optimization.stage'''      ,				-1		)

        # offload
        lowerCAmelCase__				:  Dict		       =     False
        if self.is_zeroa() or self.is_zeroa():
          lowerCAmelCase__				:  Dict		       =     set(['''cpu''', '''nvme''']		)
          lowerCAmelCase__				:  Optional[Any]		       =     set(
              [
                  self.get_value('''zero_optimization.offload_optimizer.device'''		),
                  self.get_value('''zero_optimization.offload_param.device'''		),
              ]		)
          if len(offload_devices & offload_devices_valid		) > 0:
            lowerCAmelCase__				:  List[Any]		       =     True

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  int		       =     self.config

        # find the config node of interest if it exists
        lowerCAmelCase__				:  Tuple		       =     ds_key_long.split('''.'''		)
        lowerCAmelCase__				:  int		       =     nodes.pop()
        for node in nodes:
          lowerCAmelCase__				:  int		       =     config.get(__UpperCAmelCase		)
          if config is None:
            return None, ds_key

        return config, ds_key

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=None		):
        lowerCAmelCase__       ,			lowerCAmelCase__				:  Optional[int]		       =     self.find_config_node(__UpperCAmelCase		)
        if config is None:
          return default
        return config.get(__UpperCAmelCase      ,				__UpperCAmelCase		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=False		):
        lowerCAmelCase__				:  Dict		       =     self.config

        # find the config node of interest if it exists
        lowerCAmelCase__				:  str		       =     ds_key_long.split('''.'''		)
        for node in nodes:
          lowerCAmelCase__				:  Dict		       =     config
          lowerCAmelCase__				:  List[str]		       =     config.get(__UpperCAmelCase		)
          if config is None:
            if must_exist:
              raise ValueError(f"""Can't find {ds_key_long} entry in the config: {self.config}"""		)
            else:
              return

        # if found remove it
        if parent_config is not None:
          parent_config.pop(__UpperCAmelCase		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Union[str, Any]		       =     self.get_value(__UpperCAmelCase		)
        return False if value is None else bool(__UpperCAmelCase		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Optional[Any]		       =     self.get_value(__UpperCAmelCase		)
        return False if value is None else not bool(__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        return self._stage == 2

      def 	__magic_name__(   self		):
        return self._stage == 3







      def 	__magic_name__(   self		):
        return self._offload






class  _lowerCAmelCase     :

      def __init__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  int		       =     engine







      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				**__UpperCAmelCase		):
        # runs backpropagation and handles mixed precision
        self.engine.backward(__UpperCAmelCase      ,				**__UpperCAmelCase		)

        # Deepspeed's `engine.step` performs the following operations:
        # - gradient accumulation check
        # - gradient clipping
        # - optimizer step
        # - zero grad
        # - checking overflow
        # - lr_scheduler step (only if engine.lr_scheduler is not None)
        self.engine.step()
        # and this plugin overrides the above calls with no-ops when Accelerate runs under
        # Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
        # training loop that works transparently under many training regimes.






class  _lowerCAmelCase     (  _lowercase	):

      def __init__(   self      ,				__UpperCAmelCase		):
        super().__init__(__UpperCAmelCase      ,				device_placement=__UpperCAmelCase      ,				scaler=__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     hasattr(self.optimizer      ,				'''overflow'''		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase=None		):
        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed

      def 	__magic_name__(   self		):
        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







      @property
      def 	__magic_name__(   self		):
        if self.__has_overflow__:
          return self.optimizer.overflow
        return False






class  _lowerCAmelCase     (  _lowercase	):

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        super().__init__(__UpperCAmelCase      ,				__UpperCAmelCase		)







      def 	__magic_name__(   self		):
        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed






class  _lowerCAmelCase     :

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=0.001      ,				__UpperCAmelCase=0      ,				**__UpperCAmelCase		):
        lowerCAmelCase__				:  Any		       =     params
        lowerCAmelCase__				:  str		       =     lr
        lowerCAmelCase__				:  List[str]		       =     weight_decay
        lowerCAmelCase__				:  Dict		       =     kwargs






class  _lowerCAmelCase     :

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=None      ,				__UpperCAmelCase=0      ,				**__UpperCAmelCase		):
        lowerCAmelCase__				:  Union[str, Any]		       =     optimizer
        lowerCAmelCase__				:  Optional[Any]		       =     total_num_steps
        lowerCAmelCase__				:  str		       =     warmup_num_steps
        lowerCAmelCase__				:  Dict		       =     kwargs

