





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					):
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	0.0
					lowerCAmelCase_							=	1
					lowerCAmelCase_							=	1
					lowerCAmelCase_							=	True
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	jnp.floataa




					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				[]
										_snake_case									=				[]

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

															_snake_case									=				FlaxResnetBlockaD(
															    in_channels=lowerCAmelCase_				,						out_channels=self.out_channels				,						dropout_prob=self.dropout				,						dtype=self.dtype				,						)
															resnets.append(lowerCAmelCase_  )

															_snake_case									=				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_  )

										_snake_case									=				resnets
										_snake_case									=				attentions

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






					def __call__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=True  )		->							int:
										_snake_case									=				()

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

										if self.add_downsample:
															_snake_case									=				self.downsamplers_a(lowerCAmelCase_  )
															output_states += (hidden_states,)

										return hidden_states, output_states





class 			UpperCamelCase_		(  nn.Module					):
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	0.0
					lowerCAmelCase_							=	1
					lowerCAmelCase_							=	True
					lowerCAmelCase_							=	jnp.floataa




					def 	lowerCAmelCase			(		self  )		->							Any:
										_snake_case									=				[]

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

															_snake_case									=				FlaxResnetBlockaD(
															    in_channels=lowerCAmelCase_				,						out_channels=self.out_channels				,						dropout_prob=self.dropout				,						dtype=self.dtype				,						)
															resnets.append(lowerCAmelCase_  )
										_snake_case									=				resnets

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






					def __call__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=True  )		->							Tuple:
										_snake_case									=				()

										for resnet in self.resnets:
															_snake_case									=				resnet(lowerCAmelCase_				,						lowerCAmelCase_				,						deterministic=lowerCAmelCase_  )
															output_states += (hidden_states,)

										if self.add_downsample:
															_snake_case									=				self.downsamplers_a(lowerCAmelCase_  )
															output_states += (hidden_states,)

										return hidden_states, output_states





class 			UpperCamelCase_		(  nn.Module					):
					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 	lowerCAmelCase			(		self  )		->							Union[str, Any]:
										_snake_case									=				[]
										_snake_case									=				[]

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

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

															_snake_case									=				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_  )

										_snake_case									=				resnets
										_snake_case									=				attentions

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






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

															_snake_case									=				resnet(lowerCAmelCase_				,						lowerCAmelCase_				,						deterministic=lowerCAmelCase_  )
															_snake_case									=				attn(lowerCAmelCase_				,						lowerCAmelCase_				,						deterministic=lowerCAmelCase_  )

										if self.add_upsample:
															_snake_case									=				self.upsamplers_a(lowerCAmelCase_  )

										return hidden_states





class 			UpperCamelCase_		(  nn.Module					):
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	0.0
					lowerCAmelCase_							=	1
					lowerCAmelCase_							=	True
					lowerCAmelCase_							=	jnp.floataa




					def 	lowerCAmelCase			(		self  )		->							Union[str, Any]:
										_snake_case									=				[]

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

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

										_snake_case									=				resnets

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






					def __call__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=True  )		->							List[Any]:
										for resnet in self.resnets:
															# pop res hidden states
															_snake_case									=				res_hidden_states_tuple[-1]
															_snake_case									=				res_hidden_states_tuple[:-1]
															_snake_case									=				jnp.concatenate((hidden_states, res_hidden_states)				,						axis=-1  )

															_snake_case									=				resnet(lowerCAmelCase_				,						lowerCAmelCase_				,						deterministic=lowerCAmelCase_  )

										if self.add_upsample:
															_snake_case									=				self.upsamplers_a(lowerCAmelCase_  )

										return hidden_states





class 			UpperCamelCase_		(  nn.Module					):
					lowerCAmelCase_							=	42
					lowerCAmelCase_							=	0.0
					lowerCAmelCase_							=	1
					lowerCAmelCase_							=	1
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	jnp.floataa




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

										_snake_case									=				[]

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

															_snake_case									=				FlaxResnetBlockaD(
															    in_channels=self.in_channels				,						out_channels=self.in_channels				,						dropout_prob=self.dropout				,						dtype=self.dtype				,						)
															resnets.append(lowerCAmelCase_  )

										_snake_case									=				resnets
										_snake_case									=				attentions






					def __call__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=True  )		->							Tuple:
										_snake_case									=				self.resnets[0](lowerCAmelCase_				,						lowerCAmelCase_  )
										for attn, resnet in zip(self.attentions				,						self.resnets[1:]  ):
															_snake_case									=				attn(lowerCAmelCase_				,						lowerCAmelCase_				,						deterministic=lowerCAmelCase_  )
															_snake_case									=				resnet(lowerCAmelCase_				,						lowerCAmelCase_				,						deterministic=lowerCAmelCase_  )

										return hidden_states










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_      							=     {
    """configuration_bigbird_pegasus""": [
        """BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """BigBirdPegasusConfig""",
        """BigBirdPegasusOnnxConfig""",
    ],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase_      							=     [
					    """BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """BigBirdPegasusForCausalLM""",
					    """BigBirdPegasusForConditionalGeneration""",
					    """BigBirdPegasusForQuestionAnswering""",
					    """BigBirdPegasusForSequenceClassification""",
					    """BigBirdPegasusModel""",
					    """BigBirdPegasusPreTrainedModel""",
					]


if TYPE_CHECKING:
					from .configuration_bigbird_pegasus import (
					    BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    BigBirdPegasusConfig,
					    BigBirdPegasusOnnxConfig,
					)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_bigbird_pegasus import (
										    BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
										    BigBirdPegasusForCausalLM,
										    BigBirdPegasusForConditionalGeneration,
										    BigBirdPegasusForQuestionAnswering,
										    BigBirdPegasusForSequenceClassification,
										    BigBirdPegasusModel,
										    BigBirdPegasusPreTrainedModel,
										)


else:
					import sys

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



