


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





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					int
    lowerCamelCase  :					float							     =			0.0
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					bool							     =			True
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     List[Any]:
         __lowerCamelCase   :	Dict								=		[]
         __lowerCamelCase   :	Dict								=		[]

         for i in range(self.num_layers	):
              __lowerCamelCase   :	List[Any]								=		self.in_channels if i == 0 else self.out_channels

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

              __lowerCamelCase   :	List[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(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Dict								=		resnets
         __lowerCamelCase   :	List[Any]								=		attentions

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







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=True	)				->     Union[str, Any]:
         __lowerCamelCase   :	str								=		()

         for resnet, attn in zip(self.resnets			,     self.attentions	):
              __lowerCamelCase   :	Any								=		resnet(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	Optional[int]								=		attn(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)
              output_states += (hidden_states,)

         if self.add_downsample:
              __lowerCamelCase   :	List[Any]								=		self.downsamplers_a(SCREAMING_SNAKE_CASE_	)
              output_states += (hidden_states,)

         return hidden_states, output_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					int
    lowerCamelCase  :					float							     =			0.0
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					bool							     =			True
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     str:
         __lowerCamelCase   :	int								=		[]

         for i in range(self.num_layers	):
              __lowerCamelCase   :	List[str]								=		self.in_channels if i == 0 else self.out_channels

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

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







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=True	)				->     Tuple:
         __lowerCamelCase   :	int								=		()

         for resnet in self.resnets:
              __lowerCamelCase   :	Any								=		resnet(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)
              output_states += (hidden_states,)

         if self.add_downsample:
              __lowerCamelCase   :	Optional[int]								=		self.downsamplers_a(SCREAMING_SNAKE_CASE_	)
              output_states += (hidden_states,)

         return hidden_states, output_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					int
    lowerCamelCase  :					int
    lowerCamelCase  :					float							     =			0.0
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					bool							     =			True
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     Optional[Any]:
         __lowerCamelCase   :	str								=		[]
         __lowerCamelCase   :	Dict								=		[]

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

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

              __lowerCamelCase   :	Optional[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(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	str								=		resnets
         __lowerCamelCase   :	str								=		attentions

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







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=True	)				->     Dict:
         for resnet, attn in zip(self.resnets			,     self.attentions	):
              # pop res hidden states
              __lowerCamelCase   :	str								=		res_hidden_states_tuple[-1]
              __lowerCamelCase   :	Dict								=		res_hidden_states_tuple[:-1]
              __lowerCamelCase   :	Dict								=		jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1	)

              __lowerCamelCase   :	Tuple								=		resnet(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	str								=		attn(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)

         if self.add_upsample:
              __lowerCamelCase   :	Tuple								=		self.upsamplers_a(SCREAMING_SNAKE_CASE_	)

         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					int
    lowerCamelCase  :					int
    lowerCamelCase  :					float							     =			0.0
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					bool							     =			True
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	Optional[Any]								=		[]

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

              __lowerCamelCase   :	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(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	List[str]								=		resnets

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







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=True	)				->     Optional[int]:
         for resnet in self.resnets:
              # pop res hidden states
              __lowerCamelCase   :	str								=		res_hidden_states_tuple[-1]
              __lowerCamelCase   :	int								=		res_hidden_states_tuple[:-1]
              __lowerCamelCase   :	Tuple								=		jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1	)

              __lowerCamelCase   :	List[str]								=		resnet(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)

         if self.add_upsample:
              __lowerCamelCase   :	Optional[int]								=		self.upsamplers_a(SCREAMING_SNAKE_CASE_	)

         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    lowerCamelCase  :					int
    lowerCamelCase  :					float							     =			0.0
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					int							     =			1
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					bool							     =			False
    lowerCamelCase  :					jnp.dtype							     =			jnp.floataa







    def       lowercase_				(							self	)				->     Any:
         # there is always at least one resnet
         __lowerCamelCase   :	Optional[Any]								=		[
             FlaxResnetBlockaD(
                 in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
         ]

         __lowerCamelCase   :	Tuple								=		[]

         for _ in range(self.num_layers	):
              __lowerCamelCase   :	Dict								=		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(SCREAMING_SNAKE_CASE_	)

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

         __lowerCamelCase   :	List[str]								=		resnets
         __lowerCamelCase   :	Optional[int]								=		attentions







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=True	)				->     Optional[Any]:
         __lowerCamelCase   :	List[Any]								=		self.resnets[0](SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         for attn, resnet in zip(self.attentions			,     self.resnets[1:]	):
              __lowerCamelCase   :	Tuple								=		attn(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	List[Any]								=		resnet(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     deterministic=SCREAMING_SNAKE_CASE_	)

         return hidden_states










'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A__		:					List[Any]												=  {
    """configuration_megatron_bert""": ["""MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MegatronBertConfig"""],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__		:					int												=  [
          """MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MegatronBertForCausalLM""",
          """MegatronBertForMaskedLM""",
          """MegatronBertForMultipleChoice""",
          """MegatronBertForNextSentencePrediction""",
          """MegatronBertForPreTraining""",
          """MegatronBertForQuestionAnswering""",
          """MegatronBertForSequenceClassification""",
          """MegatronBertForTokenClassification""",
          """MegatronBertModel""",
          """MegatronBertPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_megatron_bert import (
                MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                MegatronBertForCausalLM,
                MegatronBertForMaskedLM,
                MegatronBertForMultipleChoice,
                MegatronBertForNextSentencePrediction,
                MegatronBertForPreTraining,
                MegatronBertForQuestionAnswering,
                MegatronBertForSequenceClassification,
                MegatronBertForTokenClassification,
                MegatronBertModel,
                MegatronBertPreTrainedModel,
            )

else:
      import sys

      A__		:					List[Any]												=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






