


'''simple docstring'''



import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


_lowercase      :      Union[str, Any]            =					logging.get_logger(__name__)

# General docstring
_lowercase      :      Optional[int]            =					"""PoolFormerConfig"""

# Base docstring
_lowercase      :      str            =					"""sail/poolformer_s12"""
_lowercase      :      List[str]            =					[1, 512, 7, 7]

# Image classification docstring
_lowercase      :      List[str]            =					"""sail/poolformer_s12"""
_lowercase      :      List[Any]            =					"""tabby, tabby cat"""

_lowercase      :      str            =					[
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]


def 				lowerCamelCase__				(	A       :      List[Any]				,       A       :      str = 0.0				,       A       :      Any = False					):




      '''simple docstring'''






      if drop_prob == 0.0 or not training:
            return input
      UpperCAmelCase					    =     1 - drop_prob
      UpperCAmelCase					    =     (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
      UpperCAmelCase					    =     keep_prob + torch.rand(a__				,       dtype=input.dtype				,       device=input.device					)
      random_tensor.floor_()  # binarize
      UpperCAmelCase					    =     input.div(a__					) * random_tensor
      return output






class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			int				,   lowerCAmelCase      :			Optional[float] = None				)->					None:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     drop_prob




    def 					a__(		self      :			Any				,   lowerCAmelCase      :			torch.Tensor				)->					torch.Tensor:

          """simple docstring"""




          return drop_path(lowercase_				,   self.drop_prob				,   self.training				)





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

          """simple docstring"""




          return "p={}".format(self.drop_prob				)







class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			Tuple				,   lowerCAmelCase      :			List[Any]				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			List[str]=None				)->					Optional[Any]:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     patch_size if isinstance(lowercase_				,   collections.abc.Iterable				) else (patch_size, patch_size)
          UpperCAmelCase					    =     stride if isinstance(lowercase_				,   collections.abc.Iterable				) else (stride, stride)
          UpperCAmelCase					    =     padding if isinstance(lowercase_				,   collections.abc.Iterable				) else (padding, padding)

          UpperCAmelCase					    =     nn.Convad(lowercase_				,   lowercase_				,   kernel_size=lowercase_				,   stride=lowercase_				,   padding=lowercase_				)
          UpperCAmelCase					    =     norm_layer(lowercase_				) if norm_layer else nn.Identity()





    def 					a__(		self      :			Dict				,   lowerCAmelCase      :			int				)->					str:

          """simple docstring"""




          UpperCAmelCase					    =     self.projection(lowercase_				)
          UpperCAmelCase					    =     self.norm(lowercase_				)
          return embeddings







class 	UpperCamelCase__(       nn.GroupNorm	):




    def __init__(		self      :			Any				,   lowerCAmelCase      :			str				,   **lowerCAmelCase      :			Dict				)->					Optional[Any]:

          """simple docstring"""




          super().__init__(1				,   lowercase_				,   **lowercase_				)







class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			List[str]				,   lowerCAmelCase      :			Any				)->					int:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     nn.AvgPoolad(lowercase_				,   stride=1				,   padding=pool_size // 2				,   count_include_pad=lowercase_				)





    def 					a__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			Optional[Any]				)->					int:

          """simple docstring"""




          return self.pool(lowercase_				) - hidden_states







class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			str				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Optional[int]				)->					List[Any]:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     nn.Convad(lowercase_				,   lowercase_				,   1				)
          UpperCAmelCase					    =     nn.Convad(lowercase_				,   lowercase_				,   1				)
          UpperCAmelCase					    =     PoolFormerDropPath(lowercase_				)
          if isinstance(config.hidden_act				,   lowercase_				):
                UpperCAmelCase					    =     ACTaFN[config.hidden_act]
          else:
                UpperCAmelCase					    =     config.hidden_act





    def 					a__(		self      :			Optional[int]				,   lowerCAmelCase      :			Any				)->					List[str]:

          """simple docstring"""




          UpperCAmelCase					    =     self.conva(lowercase_				)
          UpperCAmelCase					    =     self.act_fn(lowercase_				)
          UpperCAmelCase					    =     self.drop(lowercase_				)
          UpperCAmelCase					    =     self.conva(lowercase_				)
          UpperCAmelCase					    =     self.drop(lowercase_				)

          return hidden_states







class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			int				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			List[Any]				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			int				)->					List[Any]:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     PoolFormerPooling(lowercase_				)
          UpperCAmelCase					    =     PoolFormerOutput(lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				)
          UpperCAmelCase					    =     PoolFormerGroupNorm(lowercase_				)
          UpperCAmelCase					    =     PoolFormerGroupNorm(lowercase_				)

          # Useful for training neural nets
          UpperCAmelCase					    =     PoolFormerDropPath(lowercase_				) if drop_path > 0.0 else nn.Identity()
          UpperCAmelCase					    =     config.use_layer_scale
          if config.use_layer_scale:
                UpperCAmelCase					    =     nn.Parameter(
                    config.layer_scale_init_value * torch.ones((lowercase_)				)				,   requires_grad=lowercase_				)
                UpperCAmelCase					    =     nn.Parameter(
                    config.layer_scale_init_value * torch.ones((lowercase_)				)				,   requires_grad=lowercase_				)





    def 					a__(		self      :			Optional[int]				,   lowerCAmelCase      :			Any				)->					int:

          """simple docstring"""




          if self.use_layer_scale:
                UpperCAmelCase					    =     self.pooling(self.before_norm(lowercase_				)				)
                UpperCAmelCase					    =     self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * pooling_output
                # First residual connection
                UpperCAmelCase					    =     hidden_states + self.drop_path(lowercase_				)
                UpperCAmelCase					    =     ()

                UpperCAmelCase					    =     self.output(self.after_norm(lowercase_				)				)
                UpperCAmelCase					    =     self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * layer_output
                # Second residual connection
                UpperCAmelCase					    =     hidden_states + self.drop_path(lowercase_				)

                UpperCAmelCase					    =     (output,) + outputs
                return outputs

          else:
                UpperCAmelCase					    =     self.drop_path(self.pooling(self.before_norm(lowercase_				)				)				)
                # First residual connection
                UpperCAmelCase					    =     pooling_output + hidden_states
                UpperCAmelCase					    =     ()

                # Second residual connection inside the PoolFormerOutput block
                UpperCAmelCase					    =     self.drop_path(self.output(self.after_norm(lowercase_				)				)				)
                UpperCAmelCase					    =     hidden_states + layer_output

                UpperCAmelCase					    =     (output,) + outputs
                return outputs







class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			int				,   lowerCAmelCase      :			Optional[int]				)->					Any:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     config
          # stochastic depth decay rule
          UpperCAmelCase					    =     [x.item() for x in torch.linspace(0				,   config.drop_path_rate				,   sum(config.depths				)				)]

          # patch embeddings
          UpperCAmelCase					    =     []
          for i in range(config.num_encoder_blocks				):
                embeddings.append(
                    PoolFormerEmbeddings(
                        patch_size=config.patch_sizes[i]				,   stride=config.strides[i]				,   padding=config.padding[i]				,   num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]				,   hidden_size=config.hidden_sizes[i]				,   )				)
          UpperCAmelCase					    =     nn.ModuleList(lowercase_				)

          # Transformer blocks
          UpperCAmelCase					    =     []
          UpperCAmelCase					    =     0
          for i in range(config.num_encoder_blocks				):
                # each block consists of layers
                UpperCAmelCase					    =     []
                if i != 0:
                      cur += config.depths[i - 1]
                for j in range(config.depths[i]				):
                      layers.append(
                          PoolFormerLayer(
                              lowercase_				,   num_channels=config.hidden_sizes[i]				,   pool_size=config.pool_size				,   hidden_size=config.hidden_sizes[i]				,   intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio				)				,   drop_path=dpr[cur + j]				,   )				)
                blocks.append(nn.ModuleList(lowercase_				)				)

          UpperCAmelCase					    =     nn.ModuleList(lowercase_				)





    def 					a__(		self      :			Dict				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			List[Any]=False				,   lowerCAmelCase      :			Dict=True				)->					Optional[Any]:

          """simple docstring"""




          UpperCAmelCase					    =     () if output_hidden_states else None

          UpperCAmelCase					    =     pixel_values
          for idx, layers in enumerate(zip(self.patch_embeddings				,   self.block				)				):
                UpperCAmelCase    ,    UpperCAmelCase					    =     layers
                # Get patch embeddings from hidden_states
                UpperCAmelCase					    =     embedding_layer(lowercase_				)
                # Send the embeddings through the blocks
                for _, blk in enumerate(lowercase_				):
                      UpperCAmelCase					    =     blk(lowercase_				)
                      UpperCAmelCase					    =     layer_outputs[0]

                if output_hidden_states:
                      UpperCAmelCase					    =     all_hidden_states + (hidden_states,)

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

          return BaseModelOutputWithNoAttention(last_hidden_state=lowercase_				,   hidden_states=lowercase_				)







class 	UpperCamelCase__(       lowerCAmelCase	):
    __magic_name__    :				List[str]  		=	PoolFormerConfig
    __magic_name__    :				Optional[Any]  		=	"poolformer"
    __magic_name__    :				Union[str, Any]  		=	"pixel_values"
    __magic_name__    :				Optional[Any]  		=	True




    def 					a__(		self      :			List[str]				,   lowerCAmelCase      :			Optional[Any]				)->					Optional[Any]:

          """simple docstring"""




          if isinstance(lowercase_				,   (nn.Linear, nn.Convad)				):
                module.weight.data.normal_(mean=0.0				,   std=self.config.initializer_range				)
                if module.bias is not None:
                      module.bias.data.zero_()
          elif isinstance(lowercase_				,   nn.LayerNorm				):
                module.bias.data.zero_()
                module.weight.data.fill_(1.0				)





    def 					a__(		self      :			Dict				,   lowerCAmelCase      :			str				,   lowerCAmelCase      :			List[Any]=False				)->					int:

          """simple docstring"""




          if isinstance(lowercase_				,   lowercase_				):
                UpperCAmelCase					    =     value







_lowercase      :      Tuple            =					r"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): 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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_lowercase      :      Optional[Any]            =					r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""






@add_start_docstrings(
    "The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."						,    lowerCAmelCase						,    )
class 	UpperCamelCase__(       lowerCAmelCase	):




    def __init__(		self      :			Optional[int]				,   lowerCAmelCase      :			Any				)->					Union[str, Any]:

          """simple docstring"""




          super().__init__(lowercase_				)
          UpperCAmelCase					    =     config

          UpperCAmelCase					    =     PoolFormerEncoder(lowercase_				)

          # Initialize weights and apply final processing
          self.post_init()




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

          """simple docstring"""




          return self.embeddings.patch_embeddings





    @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 					a__(		self      :			Tuple				,   lowerCAmelCase      :			Optional[torch.FloatTensor] = None				,   lowerCAmelCase      :			Optional[bool] = None				,   lowerCAmelCase      :			Optional[bool] = None				,   )->					Union[Tuple, BaseModelOutputWithNoAttention]:

          """simple docstring"""




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

          if pixel_values is None:
                raise ValueError('''You have to specify pixel_values'''				)

          UpperCAmelCase					    =     self.encoder(
              lowercase_				,   output_hidden_states=lowercase_				,   return_dict=lowercase_				,   )
          UpperCAmelCase					    =     encoder_outputs[0]

          if not return_dict:
                return (sequence_output, None) + encoder_outputs[1:]

          return BaseModelOutputWithNoAttention(
              last_hidden_state=lowercase_				,   hidden_states=encoder_outputs.hidden_states				,   )







class 	UpperCamelCase__(       nn.Module	):




    def __init__(		self      :			List[Any]				,   lowerCAmelCase      :			List[Any]				)->					Dict:

          """simple docstring"""




          super().__init__()
          UpperCAmelCase					    =     nn.Linear(config.hidden_size				,   config.hidden_size				)





    def 					a__(		self      :			List[Any]				,   lowerCAmelCase      :			str				)->					Union[str, Any]:

          """simple docstring"""




          UpperCAmelCase					    =     self.dense(lowercase_				)
          return output







@add_start_docstrings(
    "\n    PoolFormer Model transformer with an image classification head on top\n    "						,    lowerCAmelCase						,    )
class 	UpperCamelCase__(       lowerCAmelCase	):




    def __init__(		self      :			Tuple				,   lowerCAmelCase      :			str				)->					int:

          """simple docstring"""




          super().__init__(lowercase_				)
          UpperCAmelCase					    =     config.num_labels
          UpperCAmelCase					    =     PoolFormerModel(lowercase_				)

          # Final norm
          UpperCAmelCase					    =     PoolFormerGroupNorm(config.hidden_sizes[-1]				)
          # Classifier head
          UpperCAmelCase					    =     (
              nn.Linear(config.hidden_sizes[-1]				,   config.num_labels				) if config.num_labels > 0 else nn.Identity()
          )

          # Initialize weights and apply final processing
          self.post_init()





    @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 					a__(		self      :			Dict				,   lowerCAmelCase      :			Optional[torch.FloatTensor] = None				,   lowerCAmelCase      :			Optional[torch.LongTensor] = None				,   lowerCAmelCase      :			Optional[bool] = None				,   lowerCAmelCase      :			Optional[bool] = None				,   )->					Union[Tuple, ImageClassifierOutputWithNoAttention]:

          """simple docstring"""




          UpperCAmelCase					    =     return_dict if return_dict is not None else self.config.use_return_dict

          UpperCAmelCase					    =     self.poolformer(
              lowercase_				,   output_hidden_states=lowercase_				,   return_dict=lowercase_				,   )

          UpperCAmelCase					    =     outputs[0]

          UpperCAmelCase					    =     self.classifier(self.norm(lowercase_				).mean([-2, -1]				)				)

          UpperCAmelCase					    =     None
          if labels is not None:
                if self.config.problem_type is None:
                      if self.num_labels == 1:
                            UpperCAmelCase					    =     '''regression'''
                      elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                            UpperCAmelCase					    =     '''single_label_classification'''
                      else:
                            UpperCAmelCase					    =     '''multi_label_classification'''

                if self.config.problem_type == "regression":
                      UpperCAmelCase					    =     MSELoss()
                      if self.num_labels == 1:
                            UpperCAmelCase					    =     loss_fct(logits.squeeze()				,   labels.squeeze()				)
                      else:
                            UpperCAmelCase					    =     loss_fct(lowercase_				,   lowercase_				)
                elif self.config.problem_type == "single_label_classification":
                      UpperCAmelCase					    =     CrossEntropyLoss()
                      UpperCAmelCase					    =     loss_fct(logits.view(-1				,   self.num_labels				)				,   labels.view(-1				)				)
                elif self.config.problem_type == "multi_label_classification":
                      UpperCAmelCase					    =     BCEWithLogitsLoss()
                      UpperCAmelCase					    =     loss_fct(lowercase_				,   lowercase_				)

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

          return ImageClassifierOutputWithNoAttention(loss=lowercase_				,   logits=lowercase_				,   hidden_states=outputs.hidden_states				)









'''simple docstring'''



import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class 	UpperCamelCase__(       nn.Module	):
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				float  		=	0.0
				__magic_name__    :				int  		=	1
				__magic_name__    :				int  		=	1
				__magic_name__    :				bool  		=	True
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				jnp.dtype  		=	jnp.floataa




				def 					a__(		self      :			str				)->					Dict:

										"""simple docstring"""




										UpperCAmelCase					    =     []
										UpperCAmelCase					    =     []

										for i in range(self.num_layers				):
																UpperCAmelCase					    =     self.in_channels if i == 0 else self.out_channels

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=lowerCAmelCase				,   out_channels=self.out_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)

																UpperCAmelCase					    =     FlaxTransformeraDModel(
																    in_channels=self.out_channels				,   n_heads=self.num_attention_heads				,   d_head=self.out_channels // self.num_attention_heads				,   depth=1				,   use_linear_projection=self.use_linear_projection				,   only_cross_attention=self.only_cross_attention				,   use_memory_efficient_attention=self.use_memory_efficient_attention				,   dtype=self.dtype				,   )
																attentions.append(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets
										UpperCAmelCase					    =     attentions

										if self.add_downsample:
																UpperCAmelCase					    =     FlaxDownsampleaD(self.out_channels				,   dtype=self.dtype				)





				def __call__(		self      :			List[Any]				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			List[str]=True				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     ()

										for resnet, attn in zip(self.resnets				,   self.attentions				):
																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																UpperCAmelCase					    =     attn(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																output_states += (hidden_states,)

										if self.add_downsample:
																UpperCAmelCase					    =     self.downsamplers_a(lowerCAmelCase				)
																output_states += (hidden_states,)

										return hidden_states, output_states







class 	UpperCamelCase__(       nn.Module	):
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				float  		=	0.0
				__magic_name__    :				int  		=	1
				__magic_name__    :				bool  		=	True
				__magic_name__    :				jnp.dtype  		=	jnp.floataa




				def 					a__(		self      :			List[str]				)->					Any:

										"""simple docstring"""




										UpperCAmelCase					    =     []

										for i in range(self.num_layers				):
																UpperCAmelCase					    =     self.in_channels if i == 0 else self.out_channels

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=lowerCAmelCase				,   out_channels=self.out_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)
										UpperCAmelCase					    =     resnets

										if self.add_downsample:
																UpperCAmelCase					    =     FlaxDownsampleaD(self.out_channels				,   dtype=self.dtype				)





				def __call__(		self      :			List[str]				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			List[str]=True				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     ()

										for resnet in self.resnets:
																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																output_states += (hidden_states,)

										if self.add_downsample:
																UpperCAmelCase					    =     self.downsamplers_a(lowerCAmelCase				)
																output_states += (hidden_states,)

										return hidden_states, output_states







class 	UpperCamelCase__(       nn.Module	):
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				float  		=	0.0
				__magic_name__    :				int  		=	1
				__magic_name__    :				int  		=	1
				__magic_name__    :				bool  		=	True
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				jnp.dtype  		=	jnp.floataa




				def 					a__(		self      :			List[str]				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     []
										UpperCAmelCase					    =     []

										for i in range(self.num_layers				):
																UpperCAmelCase					    =     self.in_channels if (i == self.num_layers - 1) else self.out_channels
																UpperCAmelCase					    =     self.prev_output_channel if i == 0 else self.out_channels

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=resnet_in_channels + res_skip_channels				,   out_channels=self.out_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)

																UpperCAmelCase					    =     FlaxTransformeraDModel(
																    in_channels=self.out_channels				,   n_heads=self.num_attention_heads				,   d_head=self.out_channels // self.num_attention_heads				,   depth=1				,   use_linear_projection=self.use_linear_projection				,   only_cross_attention=self.only_cross_attention				,   use_memory_efficient_attention=self.use_memory_efficient_attention				,   dtype=self.dtype				,   )
																attentions.append(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets
										UpperCAmelCase					    =     attentions

										if self.add_upsample:
																UpperCAmelCase					    =     FlaxUpsampleaD(self.out_channels				,   dtype=self.dtype				)





				def __call__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			Union[str, Any]=True				)->					Optional[int]:

										"""simple docstring"""




										for resnet, attn in zip(self.resnets				,   self.attentions				):
																# pop res hidden states
																UpperCAmelCase					    =     res_hidden_states_tuple[-1]
																UpperCAmelCase					    =     res_hidden_states_tuple[:-1]
																UpperCAmelCase					    =     jnp.concatenate((hidden_states, res_hidden_states)				,   axis=-1				)

																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																UpperCAmelCase					    =     attn(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)

										if self.add_upsample:
																UpperCAmelCase					    =     self.upsamplers_a(lowerCAmelCase				)

										return hidden_states







class 	UpperCamelCase__(       nn.Module	):
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				float  		=	0.0
				__magic_name__    :				int  		=	1
				__magic_name__    :				bool  		=	True
				__magic_name__    :				jnp.dtype  		=	jnp.floataa




				def 					a__(		self      :			Optional[int]				)->					str:

										"""simple docstring"""




										UpperCAmelCase					    =     []

										for i in range(self.num_layers				):
																UpperCAmelCase					    =     self.in_channels if (i == self.num_layers - 1) else self.out_channels
																UpperCAmelCase					    =     self.prev_output_channel if i == 0 else self.out_channels

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=resnet_in_channels + res_skip_channels				,   out_channels=self.out_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets

										if self.add_upsample:
																UpperCAmelCase					    =     FlaxUpsampleaD(self.out_channels				,   dtype=self.dtype				)





				def __call__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Dict=True				)->					Tuple:

										"""simple docstring"""




										for resnet in self.resnets:
																# pop res hidden states
																UpperCAmelCase					    =     res_hidden_states_tuple[-1]
																UpperCAmelCase					    =     res_hidden_states_tuple[:-1]
																UpperCAmelCase					    =     jnp.concatenate((hidden_states, res_hidden_states)				,   axis=-1				)

																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)

										if self.add_upsample:
																UpperCAmelCase					    =     self.upsamplers_a(lowerCAmelCase				)

										return hidden_states







class 	UpperCamelCase__(       nn.Module	):
				__magic_name__    :				int
				__magic_name__    :				float  		=	0.0
				__magic_name__    :				int  		=	1
				__magic_name__    :				int  		=	1
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				jnp.dtype  		=	jnp.floataa




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

										"""simple docstring"""




										UpperCAmelCase					    =     [
										    FlaxResnetBlockaD(
										        in_channels=self.in_channels				,   out_channels=self.in_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
										]

										UpperCAmelCase					    =     []

										for _ in range(self.num_layers				):
																UpperCAmelCase					    =     FlaxTransformeraDModel(
																    in_channels=self.in_channels				,   n_heads=self.num_attention_heads				,   d_head=self.in_channels // self.num_attention_heads				,   depth=1				,   use_linear_projection=self.use_linear_projection				,   use_memory_efficient_attention=self.use_memory_efficient_attention				,   dtype=self.dtype				,   )
																attentions.append(lowerCAmelCase				)

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=self.in_channels				,   out_channels=self.in_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets
										UpperCAmelCase					    =     attentions





				def __call__(		self      :			Dict				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			Any=True				)->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     self.resnets[0](lowerCAmelCase				,   lowerCAmelCase				)
										for attn, resnet in zip(self.attentions				,   self.resnets[1:]				):
																UpperCAmelCase					    =     attn(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)

										return hidden_states





