


'''simple docstring'''



import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
      from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


__UpperCAmelCase    :int				       =     logging.get_logger(__name__)  # pylint: disable=invalid-name

__UpperCAmelCase    :List[Any]				       =     2_5_6



class 			a       (	_a ):


     """simple docstring"""






     SCREAMING_SNAKE_CASE :				Dict							     =     ["melgan"]

     def __init__(       self  :	str							,				snake_case  :	SpectrogramNotesEncoder							,				snake_case  :	SpectrogramContEncoder							,				snake_case  :	TaFilmDecoder							,				snake_case  :	DDPMScheduler							,				snake_case  :	OnnxRuntimeModel if is_onnx_available() else Any							,				)		->							None:
          super().__init__()

          # From MELGAN
          __UpperCAmelCase			:				Tuple													=    math.log(1E-5    )  # Matches MelGAN training.
          __UpperCAmelCase			:				Dict													=    4.0  # Largest value for most examples
          __UpperCAmelCase			:				int													=    128

          self.register_modules(
              notes_encoder=snake_case							,				continuous_encoder=snake_case							,				decoder=snake_case							,				scheduler=snake_case							,				melgan=snake_case							,				)

     def 				lowerCamelCase__				(       self  :	Dict							,				snake_case  :	Dict							,				snake_case  :	str=(-1.0, 1.0)							,				snake_case  :	str=False    )		->							Optional[Any]:
          __UpperCAmelCase							,     __UpperCAmelCase			:				Optional[Any]													=    output_range
          if clip:
               __UpperCAmelCase			:				Dict													=    torch.clip(snake_case							,				self.min_value							,				self.max_value    )
          # Scale to [0, 1].
          __UpperCAmelCase			:				Optional[Any]													=    (features - self.min_value) / (self.max_value - self.min_value)
          # Scale to [min_out, max_out].
          return zero_one * (max_out - min_out) + min_out

     def 				lowerCamelCase__				(       self  :	Optional[Any]							,				snake_case  :	List[Any]							,				snake_case  :	str=(-1.0, 1.0)							,				snake_case  :	Union[str, Any]=False    )		->							Any:
          __UpperCAmelCase							,     __UpperCAmelCase			:				Optional[int]													=    input_range
          __UpperCAmelCase			:				Dict													=    torch.clip(snake_case							,				snake_case							,				snake_case    ) if clip else outputs
          # Scale to [0, 1].
          __UpperCAmelCase			:				Tuple													=    (outputs - min_out) / (max_out - min_out)
          # Scale to [self.min_value, self.max_value].
          return zero_one * (self.max_value - self.min_value) + self.min_value

     def 				lowerCamelCase__				(       self  :	Tuple							,				snake_case  :	Optional[int]							,				snake_case  :	Union[str, Any]							,				snake_case  :	Optional[int]    )		->							Dict:
          __UpperCAmelCase			:				Any													=    input_tokens > 0
          __UpperCAmelCase							,     __UpperCAmelCase			:				Tuple													=    self.notes_encoder(
              encoder_input_tokens=snake_case							,				encoder_inputs_mask=snake_case    )

          __UpperCAmelCase							,     __UpperCAmelCase			:				List[str]													=    self.continuous_encoder(
              encoder_inputs=snake_case							,				encoder_inputs_mask=snake_case    )

          return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]

     def 				lowerCamelCase__				(       self  :	List[str]							,				snake_case  :	Any							,				snake_case  :	List[Any]							,				snake_case  :	Optional[Any]    )		->							str:
          __UpperCAmelCase			:				List[str]													=    noise_time
          if not torch.is_tensor(snake_case    ):
               __UpperCAmelCase			:				Union[str, Any]													=    torch.tensor([timesteps]							,				dtype=torch.long							,				device=input_tokens.device    )
          elif torch.is_tensor(snake_case    ) and len(timesteps.shape    ) == 0:
               __UpperCAmelCase			:				Optional[Any]													=    timesteps[None].to(input_tokens.device    )

          # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
          __UpperCAmelCase			:				Optional[Any]													=    timesteps * torch.ones(input_tokens.shape[0]							,				dtype=timesteps.dtype							,				device=timesteps.device    )

          __UpperCAmelCase			:				Optional[int]													=    self.decoder(
              encodings_and_masks=snake_case							,				decoder_input_tokens=snake_case							,				decoder_noise_time=snake_case    )
          return logits






     @torch.no_grad()
     def __call__(       self  :	Optional[int]							,				snake_case  :	List[List[int]]							,				snake_case  :	Optional[torch.Generator] = None							,				snake_case  :	int = 100							,				snake_case  :	bool = True							,				snake_case  :	str = "numpy"							,				snake_case  :	Optional[Callable[[int, int, torch.FloatTensor], None]] = None							,				snake_case  :	int = 1							,				)		->							Union[AudioPipelineOutput, Tuple]:
          if (callback_steps is None) or (
              callback_steps is not None and (not isinstance(snake_case							,				snake_case    ) or callback_steps <= 0)
          ):
               raise ValueError(
                   f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
                   f' {type(snake_case    )}.'    )

          __UpperCAmelCase			:				Any													=    np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]							,				dtype=np.floataa    )
          __UpperCAmelCase			:				Any													=    np.zeros([1, 0, self.n_dims]							,				np.floataa    )
          __UpperCAmelCase			:				Any													=    torch.ones((1, TARGET_FEATURE_LENGTH)							,				dtype=snake_case							,				device=self.device    )

          for i, encoder_input_tokens in enumerate(snake_case    ):
               if i == 0:
                    __UpperCAmelCase			:				Optional[int]													=    torch.from_numpy(pred_mel[:1].copy()    ).to(
                        device=self.device							,				dtype=self.decoder.dtype    )
                    # The first chunk has no previous context.
                    __UpperCAmelCase			:				Optional[Any]													=    torch.zeros((1, TARGET_FEATURE_LENGTH)							,				dtype=snake_case							,				device=self.device    )
               else:
                    # The full song pipeline does not feed in a context feature, so the mask
                    # will be all 0s after the feature converter. Because we know we're
                    # feeding in a full context chunk from the previous prediction, set it
                    # to all 1s.
                    __UpperCAmelCase			:				Dict													=    ones

               __UpperCAmelCase			:				List[Any]													=    self.scale_features(
                   snake_case							,				output_range=[-1.0, 1.0]							,				clip=snake_case    )

               __UpperCAmelCase			:				List[Any]													=    self.encode(
                   input_tokens=torch.IntTensor([encoder_input_tokens]    ).to(device=self.device    )							,				continuous_inputs=snake_case							,				continuous_mask=snake_case							,				)

               # Sample encoder_continuous_inputs shaped gaussian noise to begin loop
               __UpperCAmelCase			:				List[Any]													=    randn_tensor(
                   shape=encoder_continuous_inputs.shape							,				generator=snake_case							,				device=self.device							,				dtype=self.decoder.dtype							,				)

               # set step values
               self.scheduler.set_timesteps(snake_case    )

               # Denoising diffusion loop
               for j, t in enumerate(self.progress_bar(self.scheduler.timesteps    )    ):
                    __UpperCAmelCase			:				Any													=    self.decode(
                        encodings_and_masks=snake_case							,				input_tokens=snake_case							,				noise_time=t / self.scheduler.config.num_train_timesteps							,				)

                    # Compute previous output: x_t -> x_t-1
                    __UpperCAmelCase			:				Dict													=    self.scheduler.step(snake_case							,				snake_case							,				snake_case							,				generator=snake_case    ).prev_sample

               __UpperCAmelCase			:				Dict													=    self.scale_to_features(snake_case							,				input_range=[-1.0, 1.0]    )
               __UpperCAmelCase			:				Tuple													=    mel[:1]
               __UpperCAmelCase			:				Optional[Any]													=    mel.cpu().float().numpy()

               __UpperCAmelCase			:				Dict													=    np.concatenate([full_pred_mel, pred_mel[:1]]							,				axis=1    )

               # call the callback, if provided
               if callback is not None and i % callback_steps == 0:
                    callback(snake_case							,				snake_case    )

               logger.info('''Generated segment'''							,				snake_case    )

          if output_type == "numpy" and not is_onnx_available():
               raise ValueError(
                   '''Cannot return output in \'np\' format if ONNX is not available. Make sure to have ONNX installed or set \'output_type\' to \'mel\'.'''    )
          elif output_type == "numpy" and self.melgan is None:
               raise ValueError(
                   '''Cannot return output in \'np\' format if melgan component is not defined. Make sure to define `self.melgan` or set \'output_type\' to \'mel\'.'''    )

          if output_type == "numpy":
               __UpperCAmelCase			:				Optional[int]													=    self.melgan(input_features=full_pred_mel.astype(np.floataa    )    )
          else:
               __UpperCAmelCase			:				Tuple													=    full_pred_mel

          if not return_dict:
               return (output,)

          return AudioPipelineOutput(audios=snake_case    )



'''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


__UpperCAmelCase    :Optional[Any]				       =     logging.get_logger(__name__)

# General docstring
__UpperCAmelCase    :List[Any]				       =     "RegNetConfig"

# Base docstring
__UpperCAmelCase    :List[Any]				       =     "facebook/regnet-y-040"
__UpperCAmelCase    :Union[str, Any]				       =     [1, 1_0_8_8, 7, 7]

# Image classification docstring
__UpperCAmelCase    :int				       =     "facebook/regnet-y-040"
__UpperCAmelCase    :Optional[Any]				       =     "tabby, tabby cat"

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



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


     """simple docstring"""






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






     def 				lowerCamelCase__				(       self  :	Any							,				snake_case  :	List[str]    )		->							int:
          __UpperCAmelCase			:				Tuple													=    self.convolution(self.padding(snake_case    )    )
          __UpperCAmelCase			:				List[Any]													=    self.normalization(snake_case    )
          __UpperCAmelCase			:				Optional[Any]													=    self.activation(snake_case    )
          return hidden_state



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


     """simple docstring"""






     def __init__(       self  :	str							,				snake_case  :	RegNetConfig							,				**snake_case  :	Tuple    )		->							int:
          super().__init__(**snake_case    )
          __UpperCAmelCase			:				List[str]													=    config.num_channels
          __UpperCAmelCase			:				Optional[int]													=    TFRegNetConvLayer(
              out_channels=config.embedding_size							,				kernel_size=3							,				stride=2							,				activation=config.hidden_act							,				name='''embedder'''							,				)






     def 				lowerCamelCase__				(       self  :	Optional[int]							,				snake_case  :	Dict    )		->							int:
          __UpperCAmelCase			:				int													=    shape_list(snake_case    )[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)
          __UpperCAmelCase			:				Dict													=    tf.transpose(snake_case							,				perm=(0, 2, 3, 1)    )
          __UpperCAmelCase			:				List[str]													=    self.embedder(snake_case    )
          return hidden_state



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


     """simple docstring"""






     def __init__(       self  :	Any							,				snake_case  :	int							,				snake_case  :	int = 2							,				**snake_case  :	Tuple    )		->							str:
          super().__init__(**snake_case    )
          __UpperCAmelCase			:				str													=    tf.keras.layers.ConvaD(
              filters=snake_case							,				kernel_size=1							,				strides=snake_case							,				use_bias=snake_case							,				name='''convolution'''    )
          __UpperCAmelCase			:				int													=    tf.keras.layers.BatchNormalization(epsilon=1E-5							,				momentum=0.9							,				name='''normalization'''    )






     def 				lowerCamelCase__				(       self  :	str							,				snake_case  :	tf.Tensor							,				snake_case  :	bool = False    )		->							tf.Tensor:
          return self.normalization(self.convolution(snake_case    )							,				training=snake_case    )



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


     """simple docstring"""






     def __init__(       self  :	Tuple							,				snake_case  :	int							,				snake_case  :	int							,				**snake_case  :	Tuple    )		->							List[Any]:
          super().__init__(**snake_case    )
          __UpperCAmelCase			:				List[Any]													=    tf.keras.layers.GlobalAveragePoolingaD(keepdims=snake_case							,				name='''pooler'''    )
          __UpperCAmelCase			:				Dict													=    [
              tf.keras.layers.ConvaD(filters=snake_case							,				kernel_size=1							,				activation='''relu'''							,				name='''attention.0'''    ),
              tf.keras.layers.ConvaD(filters=snake_case							,				kernel_size=1							,				activation='''sigmoid'''							,				name='''attention.2'''    ),
          ]






     def 				lowerCamelCase__				(       self  :	Optional[int]							,				snake_case  :	Tuple    )		->							Union[str, Any]:
          # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
          __UpperCAmelCase			:				str													=    self.pooler(snake_case    )
          for layer_module in self.attention:
               __UpperCAmelCase			:				int													=    layer_module(snake_case    )
          __UpperCAmelCase			:				List[Any]													=    hidden_state * pooled
          return hidden_state



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


     """simple docstring"""






     def __init__(       self  :	Any							,				snake_case  :	RegNetConfig							,				snake_case  :	int							,				snake_case  :	int							,				snake_case  :	int = 1							,				**snake_case  :	int    )		->							int:
          super().__init__(**snake_case    )
          __UpperCAmelCase			:				Any													=    in_channels != out_channels or stride != 1
          __UpperCAmelCase			:				Optional[int]													=    max(1							,				out_channels // config.groups_width    )
          __UpperCAmelCase			:				Optional[int]													=    (
              TFRegNetShortCut(snake_case							,				stride=snake_case							,				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.
          __UpperCAmelCase			:				List[Any]													=    [
              TFRegNetConvLayer(snake_case							,				kernel_size=1							,				activation=config.hidden_act							,				name='''layer.0'''    ),
              TFRegNetConvLayer(
                  snake_case							,				stride=snake_case							,				groups=snake_case							,				activation=config.hidden_act							,				name='''layer.1'''    ),
              TFRegNetConvLayer(snake_case							,				kernel_size=1							,				activation=snake_case							,				name='''layer.2'''    ),
          ]
          __UpperCAmelCase			:				Union[str, Any]													=    ACTaFN[config.hidden_act]






     def 				lowerCamelCase__				(       self  :	Union[str, Any]							,				snake_case  :	Optional[Any]    )		->							List[str]:
          __UpperCAmelCase			:				Union[str, Any]													=    hidden_state
          for layer_module in self.layers:
               __UpperCAmelCase			:				Any													=    layer_module(snake_case    )
          __UpperCAmelCase			:				Tuple													=    self.shortcut(snake_case    )
          hidden_state += residual
          __UpperCAmelCase			:				Optional[int]													=    self.activation(snake_case    )
          return hidden_state



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


     """simple docstring"""






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






     def 				lowerCamelCase__				(       self  :	Optional[Any]							,				snake_case  :	Tuple    )		->							Any:
          __UpperCAmelCase			:				Optional[int]													=    hidden_state
          for layer_module in self.layers:
               __UpperCAmelCase			:				Any													=    layer_module(snake_case    )
          __UpperCAmelCase			:				int													=    self.shortcut(snake_case    )
          hidden_state += residual
          __UpperCAmelCase			:				Optional[int]													=    self.activation(snake_case    )
          return hidden_state



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


     """simple docstring"""






     def __init__(       self  :	Optional[int]							,				snake_case  :	RegNetConfig							,				snake_case  :	int							,				snake_case  :	int							,				snake_case  :	int = 2							,				snake_case  :	int = 2							,				**snake_case  :	str    )		->							Optional[Any]:
          super().__init__(**snake_case    )

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






     def 				lowerCamelCase__				(       self  :	List[str]							,				snake_case  :	Any    )		->							List[Any]:
          for layer_module in self.layers:
               __UpperCAmelCase			:				Optional[Any]													=    layer_module(snake_case    )
          return hidden_state



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


     """simple docstring"""






     def __init__(       self  :	Any							,				snake_case  :	RegNetConfig							,				**snake_case  :	int    )		->							str:
          super().__init__(**snake_case    )
          __UpperCAmelCase			:				Dict													=    []
          # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
          self.stages.append(
              TFRegNetStage(
                  snake_case							,				config.embedding_size							,				config.hidden_sizes[0]							,				stride=2 if config.downsample_in_first_stage else 1							,				depth=config.depths[0]							,				name='''stages.0'''							,				)    )
          __UpperCAmelCase			:				Optional[Any]													=    zip(config.hidden_sizes							,				config.hidden_sizes[1:]    )
          for i, ((in_channels, out_channels), depth) in enumerate(zip(snake_case							,				config.depths[1:]    )    ):
               self.stages.append(TFRegNetStage(snake_case							,				snake_case							,				snake_case							,				depth=snake_case							,				name=f'stages.{i+1}'    )    )






     def 				lowerCamelCase__				(       self  :	int							,				snake_case  :	tf.Tensor							,				snake_case  :	bool = False							,				snake_case  :	bool = True    )		->							TFBaseModelOutputWithNoAttention:
          __UpperCAmelCase			:				Any													=    () if output_hidden_states else None

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

               __UpperCAmelCase			:				List[Any]													=    stage_module(snake_case    )

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



@keras_serializable
class 			a       (	tf.keras.layers.Layer ):


     """simple docstring"""






     SCREAMING_SNAKE_CASE :				Any							     =     RegNetConfig

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






     @unpack_inputs
     def 				lowerCamelCase__				(       self  :	Dict							,				snake_case  :	tf.Tensor							,				snake_case  :	Optional[bool] = None							,				snake_case  :	Optional[bool] = None							,				snake_case  :	bool = False							,				)		->							TFBaseModelOutputWithPoolingAndNoAttention:
          __UpperCAmelCase			:				Dict													=    (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          __UpperCAmelCase			:				Optional[Any]													=    return_dict if return_dict is not None else self.config.use_return_dict

          __UpperCAmelCase			:				Optional[Any]													=    self.embedder(snake_case							,				training=snake_case    )

          __UpperCAmelCase			:				Optional[int]													=    self.encoder(
              snake_case							,				output_hidden_states=snake_case							,				return_dict=snake_case							,				training=snake_case    )

          __UpperCAmelCase			:				List[str]													=    encoder_outputs[0]
          __UpperCAmelCase			:				str													=    self.pooler(snake_case    )

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

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



class 			a       (	_a ):


     """simple docstring"""






     SCREAMING_SNAKE_CASE :				str							     =     RegNetConfig
     SCREAMING_SNAKE_CASE :				Tuple							     =     "regnet"
     SCREAMING_SNAKE_CASE :				List[Any]							     =     "pixel_values"

     @property
     def 				lowerCamelCase__				(       self  :	int    )		->							List[str]:
          return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)							,				dtype=tf.floataa    )}



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

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



@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top."			,						_a			,						)
class 			a       (	_a ):


     """simple docstring"""






     def __init__(       self  :	List[Any]							,				snake_case  :	RegNetConfig							,				*snake_case  :	Optional[int]							,				**snake_case  :	List[str]    )		->							Tuple:
          super().__init__(snake_case							,				*snake_case							,				**snake_case    )
          __UpperCAmelCase			:				Dict													=    TFRegNetMainLayer(snake_case							,				name='''regnet'''    )






     @unpack_inputs
     @add_start_docstrings_to_model_forward(snake_case    )
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC							,				output_type=snake_case							,				config_class=_CONFIG_FOR_DOC							,				modality='''vision'''							,				expected_output=_EXPECTED_OUTPUT_SHAPE							,				)
     def 				lowerCamelCase__				(       self  :	Tuple							,				snake_case  :	tf.Tensor							,				snake_case  :	Optional[bool] = None							,				snake_case  :	Optional[bool] = None							,				snake_case  :	str=False							,				)		->							Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
          __UpperCAmelCase			:				List[Any]													=    (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          __UpperCAmelCase			:				List[str]													=    return_dict if return_dict is not None else self.config.use_return_dict

          __UpperCAmelCase			:				Dict													=    self.regnet(
              pixel_values=snake_case							,				output_hidden_states=snake_case							,				return_dict=snake_case							,				training=snake_case							,				)
          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    "			,						_a			,						)
class 			a       (	_a			,						_a ):


     """simple docstring"""






     def __init__(       self  :	Tuple							,				snake_case  :	RegNetConfig							,				*snake_case  :	Optional[Any]							,				**snake_case  :	List[Any]    )		->							List[Any]:
          super().__init__(snake_case							,				*snake_case							,				**snake_case    )
          __UpperCAmelCase			:				List[Any]													=    config.num_labels
          __UpperCAmelCase			:				Optional[int]													=    TFRegNetMainLayer(snake_case							,				name='''regnet'''    )
          # classification head
          __UpperCAmelCase			:				str													=    [
              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(snake_case    )
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT							,				output_type=snake_case							,				config_class=_CONFIG_FOR_DOC							,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,				)
     def 				lowerCamelCase__				(       self  :	Tuple							,				snake_case  :	tf.Tensor = None							,				snake_case  :	tf.Tensor = None							,				snake_case  :	bool = None							,				snake_case  :	bool = None							,				snake_case  :	Tuple=False							,				)		->							Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
          __UpperCAmelCase			:				Optional[int]													=    (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          __UpperCAmelCase			:				Optional[Any]													=    return_dict if return_dict is not None else self.config.use_return_dict

          __UpperCAmelCase			:				Optional[int]													=    self.regnet(
              snake_case							,				output_hidden_states=snake_case							,				return_dict=snake_case							,				training=snake_case    )

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

          __UpperCAmelCase			:				Tuple													=    self.classifier[0](snake_case    )
          __UpperCAmelCase			:				Tuple													=    self.classifier[1](snake_case    )

          __UpperCAmelCase			:				Any													=    None if labels is None else self.hf_compute_loss(labels=snake_case							,				logits=snake_case    )

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

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