






"""simple docstring"""







from typing import TYPE_CHECKING

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


__a       :	Union[str, Any]             = {
    'configuration_nezha': ['NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'NezhaConfig'],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__a       :	List[Any]             = [
	    'NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'NezhaForNextSentencePrediction',
	    'NezhaForMaskedLM',
	    'NezhaForPreTraining',
	    'NezhaForMultipleChoice',
	    'NezhaForQuestionAnswering',
	    'NezhaForSequenceClassification',
	    'NezhaForTokenClassification',
	    'NezhaModel',
	    'NezhaPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_nezha import NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP, NezhaConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_nezha import (
		    NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST,
		    NezhaForMaskedLM,
		    NezhaForMultipleChoice,
		    NezhaForNextSentencePrediction,
		    NezhaForPreTraining,
		    NezhaForQuestionAnswering,
		    NezhaForSequenceClassification,
		    NezhaForTokenClassification,
		    NezhaModel,
		    NezhaPreTrainedModel,
		)


else:
	import sys

	__a       :	Tuple             = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)










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





							"""simple docstring"""





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




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

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

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  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(__A    )

												a__          =  resnets
												a__          =  attentions

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





							def __call__(   self:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     str				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												a__          =  ()

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

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

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





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




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []

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

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )
												a__          =  resnets

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





							def __call__(   self:     Dict				,		__A:     int				,		__A:     Dict				,		__A:     Optional[Any]=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet in self.resnets:
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

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

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





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




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

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

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

																	a__          =  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(__A    )

												a__          =  resnets
												a__          =  attentions

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





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     List[Any]				,		__A:     List[str]				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

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

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





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




							def  lowercase		(   self:     str    ):

												'''simple docstring'''

												a__          =  []

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

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

												a__          =  resnets

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





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

												'''simple docstring'''

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

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





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




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

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

												a__          =  []

												for _ in range(self.num_layers    ):
																	a__          =  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(__A    )

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets
												a__          =  attentions





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     int				,		__A:     Tuple				,		__A:     str=True    ):

												'''simple docstring'''

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

												return hidden_states


