





"""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 		__snake_case (		nn.Module					):





 """simple docstring"""


 _lowerCamelCase				=					42
 _lowerCamelCase				=					42
 _lowerCamelCase				=					0.0
 _lowerCamelCase				=					1
 _lowerCamelCase				=					1
 _lowerCamelCase				=					True
 _lowerCamelCase				=					False
 _lowerCamelCase				=					False
 _lowerCamelCase				=					False
 _lowerCamelCase				=					jnp.floataa

 def  UpperCamelCase__(   self  ):



     '''simple docstring'''




     __A   :						int									=					[]
     __A   :						Optional[Any]									=					[]

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

         __A   :						int									=					FlaxResnetBlockaD(
             in_channels=__lowerCamelCase ,						out_channels=self.out_channels ,						dropout_prob=self.dropout ,						dtype=self.dtype ,						)
         resnets.append(__lowerCamelCase  )

         __A   :						str									=					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  )

     __A   :						Dict									=					resnets
     __A   :						int									=					attentions

     if self.add_downsample:
         __A   :						Union[str, Any]									=					FlaxDownsampleaD(self.out_channels ,						dtype=self.dtype  )







 def __call__(   self ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase=True  ):



     '''simple docstring'''




     __A   :						int									=					()

     for resnet, attn in zip(self.resnets ,						self.attentions  ):
         __A   :						Tuple									=					resnet(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )
         __A   :						List[Any]									=					attn(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )
         output_states += (hidden_states,)

     if self.add_downsample:
         __A   :						Union[str, Any]									=					self.downsamplers_a(__lowerCamelCase  )
         output_states += (hidden_states,)

     return hidden_states, output_states






class 		__snake_case (		nn.Module					):





 """simple docstring"""


 _lowerCamelCase				=					42
 _lowerCamelCase				=					42
 _lowerCamelCase				=					0.0
 _lowerCamelCase				=					1
 _lowerCamelCase				=					True
 _lowerCamelCase				=					jnp.floataa

 def  UpperCamelCase__(   self  ):



     '''simple docstring'''




     __A   :						List[str]									=					[]

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

         __A   :						Any									=					FlaxResnetBlockaD(
             in_channels=__lowerCamelCase ,						out_channels=self.out_channels ,						dropout_prob=self.dropout ,						dtype=self.dtype ,						)
         resnets.append(__lowerCamelCase  )
     __A   :						List[Any]									=					resnets

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







 def __call__(   self ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase=True  ):



     '''simple docstring'''




     __A   :						Dict									=					()

     for resnet in self.resnets:
         __A   :						Union[str, Any]									=					resnet(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )
         output_states += (hidden_states,)

     if self.add_downsample:
         __A   :						Any									=					self.downsamplers_a(__lowerCamelCase  )
         output_states += (hidden_states,)

     return hidden_states, output_states






class 		__snake_case (		nn.Module					):





 """simple docstring"""


 _lowerCamelCase				=					42
 _lowerCamelCase				=					42
 _lowerCamelCase				=					42
 _lowerCamelCase				=					0.0
 _lowerCamelCase				=					1
 _lowerCamelCase				=					1
 _lowerCamelCase				=					True
 _lowerCamelCase				=					False
 _lowerCamelCase				=					False
 _lowerCamelCase				=					False
 _lowerCamelCase				=					jnp.floataa

 def  UpperCamelCase__(   self  ):



     '''simple docstring'''




     __A   :						str									=					[]
     __A   :						Optional[int]									=					[]

     for i in range(self.num_layers  ):
         __A   :						str									=					self.in_channels if (i == self.num_layers - 1) else self.out_channels
         __A   :						Optional[Any]									=					self.prev_output_channel if i == 0 else self.out_channels

         __A   :						Optional[Any]									=					FlaxResnetBlockaD(
             in_channels=resnet_in_channels + res_skip_channels ,						out_channels=self.out_channels ,						dropout_prob=self.dropout ,						dtype=self.dtype ,						)
         resnets.append(__lowerCamelCase  )

         __A   :						Any									=					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  )

     __A   :						Any									=					resnets
     __A   :						List[Any]									=					attentions

     if self.add_upsample:
         __A   :						Union[str, Any]									=					FlaxUpsampleaD(self.out_channels ,						dtype=self.dtype  )







 def __call__(   self ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase=True  ):



     '''simple docstring'''




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

         __A   :						Optional[Any]									=					resnet(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )
         __A   :						Optional[int]									=					attn(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )

     if self.add_upsample:
         __A   :						int									=					self.upsamplers_a(__lowerCamelCase  )

     return hidden_states






class 		__snake_case (		nn.Module					):





 """simple docstring"""


 _lowerCamelCase				=					42
 _lowerCamelCase				=					42
 _lowerCamelCase				=					42
 _lowerCamelCase				=					0.0
 _lowerCamelCase				=					1
 _lowerCamelCase				=					True
 _lowerCamelCase				=					jnp.floataa

 def  UpperCamelCase__(   self  ):



     '''simple docstring'''




     __A   :						int									=					[]

     for i in range(self.num_layers  ):
         __A   :						Optional[Any]									=					self.in_channels if (i == self.num_layers - 1) else self.out_channels
         __A   :						List[str]									=					self.prev_output_channel if i == 0 else self.out_channels

         __A   :						int									=					FlaxResnetBlockaD(
             in_channels=resnet_in_channels + res_skip_channels ,						out_channels=self.out_channels ,						dropout_prob=self.dropout ,						dtype=self.dtype ,						)
         resnets.append(__lowerCamelCase  )

     __A   :						str									=					resnets

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







 def __call__(   self ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase=True  ):



     '''simple docstring'''




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

         __A   :						str									=					resnet(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )

     if self.add_upsample:
         __A   :						List[str]									=					self.upsamplers_a(__lowerCamelCase  )

     return hidden_states






class 		__snake_case (		nn.Module					):





 """simple docstring"""


 _lowerCamelCase				=					42
 _lowerCamelCase				=					0.0
 _lowerCamelCase				=					1
 _lowerCamelCase				=					1
 _lowerCamelCase				=					False
 _lowerCamelCase				=					False
 _lowerCamelCase				=					jnp.floataa

 def  UpperCamelCase__(   self  ):



     '''simple docstring'''




     __A   :						List[str]									=					[
         FlaxResnetBlockaD(
             in_channels=self.in_channels ,						out_channels=self.in_channels ,						dropout_prob=self.dropout ,						dtype=self.dtype ,						)
     ]

     __A   :						Optional[Any]									=					[]

     for _ in range(self.num_layers  ):
         __A   :						Union[str, Any]									=					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  )

         __A   :						Any									=					FlaxResnetBlockaD(
             in_channels=self.in_channels ,						out_channels=self.in_channels ,						dropout_prob=self.dropout ,						dtype=self.dtype ,						)
         resnets.append(__lowerCamelCase  )

     __A   :						List[Any]									=					resnets
     __A   :						Any									=					attentions







 def __call__(   self ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase=True  ):



     '''simple docstring'''




     __A   :						int									=					self.resnets[0](__lowerCamelCase ,						__lowerCamelCase  )
     for attn, resnet in zip(self.attentions ,						self.resnets[1:]  ):
         __A   :						Tuple									=					attn(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )
         __A   :						Optional[Any]									=					resnet(__lowerCamelCase ,						__lowerCamelCase ,						deterministic=__lowerCamelCase  )

     return hidden_states











"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


a_			=      {
    """configuration_tapas""": ["""TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TapasConfig"""],
    """tokenization_tapas""": ["""TapasTokenizer"""],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_			=      [
           """TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """TapasForMaskedLM""",
           """TapasForQuestionAnswering""",
           """TapasForSequenceClassification""",
           """TapasModel""",
           """TapasPreTrainedModel""",
           """load_tf_weights_in_tapas""",
       ]
try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_			=      [
           """TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """TFTapasForMaskedLM""",
           """TFTapasForQuestionAnswering""",
           """TFTapasForSequenceClassification""",
           """TFTapasModel""",
           """TFTapasPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
       from .tokenization_tapas import TapasTokenizer

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tapas import (
                  TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TapasForMaskedLM,
                  TapasForQuestionAnswering,
                  TapasForSequenceClassification,
                  TapasModel,
                  TapasPreTrainedModel,
                  load_tf_weights_in_tapas,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_tapas import (
                  TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFTapasForMaskedLM,
                  TFTapasForQuestionAnswering,
                  TFTapasForSequenceClassification,
                  TFTapasModel,
                  TFTapasPreTrainedModel,
              )


else:
       import sys

       a_			=      _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




