




"""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						):
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 =0.0
		UpperCamelCase						 =1
		UpperCamelCase						 =1
		UpperCamelCase						 =True
		UpperCamelCase						 =False
		UpperCamelCase						 =False
		UpperCamelCase						 =False
		UpperCamelCase						 =jnp.floataa
		def      _lowerCamelCase				(		self    )		->    Any:
					__lowercase       :   str        =    []
					__lowercase       :   Dict        =    []

					for i in range(self.num_layers    ):
								__lowercase       :   Union[str, Any]        =    self.in_channels if i == 0 else self.out_channels

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

								__lowercase       :   Dict        =    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(lowercase_    )

					__lowercase       :   Optional[int]        =    resnets
					__lowercase       :   Dict        =    attentions

					if self.add_downsample:
								__lowercase       :   Union[str, Any]        =    FlaxDownsampleaD(self.out_channels							,		dtype=self.dtype    )
		def __call__(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_=True    )		->    Union[str, Any]:
					__lowercase       :   Tuple        =    ()

					for resnet, attn in zip(self.resnets							,		self.attentions    ):
								__lowercase       :   Union[str, Any]        =    resnet(lowercase_							,		lowercase_							,		deterministic=lowercase_    )
								__lowercase       :   Union[str, Any]        =    attn(lowercase_							,		lowercase_							,		deterministic=lowercase_    )
								output_states += (hidden_states,)

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

					return hidden_states, output_states


class    UpperCAmelCase_  (       nn.Module						):
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 =0.0
		UpperCamelCase						 =1
		UpperCamelCase						 =True
		UpperCamelCase						 =jnp.floataa
		def      _lowerCamelCase				(		self    )		->    int:
					__lowercase       :   str        =    []

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

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

					if self.add_downsample:
								__lowercase       :   List[str]        =    FlaxDownsampleaD(self.out_channels							,		dtype=self.dtype    )
		def __call__(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_=True    )		->    Dict:
					__lowercase       :   Optional[int]        =    ()

					for resnet in self.resnets:
								__lowercase       :   Tuple        =    resnet(lowercase_							,		lowercase_							,		deterministic=lowercase_    )
								output_states += (hidden_states,)

					if self.add_downsample:
								__lowercase       :   Tuple        =    self.downsamplers_a(lowercase_    )
								output_states += (hidden_states,)

					return hidden_states, output_states


class    UpperCAmelCase_  (       nn.Module						):
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 =0.0
		UpperCamelCase						 =1
		UpperCamelCase						 =1
		UpperCamelCase						 =True
		UpperCamelCase						 =False
		UpperCamelCase						 =False
		UpperCamelCase						 =False
		UpperCamelCase						 =jnp.floataa
		def      _lowerCamelCase				(		self    )		->    Tuple:
					__lowercase       :   str        =    []
					__lowercase       :   List[str]        =    []

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

								__lowercase       :   List[str]        =    FlaxResnetBlockaD(
								    in_channels=resnet_in_channels + res_skip_channels							,		out_channels=self.out_channels							,		dropout_prob=self.dropout							,		dtype=self.dtype							,		)
								resnets.append(lowercase_    )

								__lowercase       :   Tuple        =    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(lowercase_    )

					__lowercase       :   int        =    resnets
					__lowercase       :   Dict        =    attentions

					if self.add_upsample:
								__lowercase       :   str        =    FlaxUpsampleaD(self.out_channels							,		dtype=self.dtype    )
		def __call__(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_=True    )		->    Optional[int]:




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

								__lowercase       :   Optional[Any]        =    resnet(lowercase_							,		lowercase_							,		deterministic=lowercase_    )
								__lowercase       :   List[Any]        =    attn(lowercase_							,		lowercase_							,		deterministic=lowercase_    )

					if self.add_upsample:
								__lowercase       :   Optional[Any]        =    self.upsamplers_a(lowercase_    )

					return hidden_states


class    UpperCAmelCase_  (       nn.Module						):
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 =0.0
		UpperCamelCase						 =1
		UpperCamelCase						 =True
		UpperCamelCase						 =jnp.floataa
		def      _lowerCamelCase				(		self    )		->    Dict:
					__lowercase       :   List[str]        =    []

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

								__lowercase       :   str        =    FlaxResnetBlockaD(
								    in_channels=resnet_in_channels + res_skip_channels							,		out_channels=self.out_channels							,		dropout_prob=self.dropout							,		dtype=self.dtype							,		)
								resnets.append(lowercase_    )

					__lowercase       :   str        =    resnets

					if self.add_upsample:
								__lowercase       :   List[str]        =    FlaxUpsampleaD(self.out_channels							,		dtype=self.dtype    )
		def __call__(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_=True    )		->    int:




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

								__lowercase       :   Optional[Any]        =    resnet(lowercase_							,		lowercase_							,		deterministic=lowercase_    )

					if self.add_upsample:
								__lowercase       :   Union[str, Any]        =    self.upsamplers_a(lowercase_    )

					return hidden_states


class    UpperCAmelCase_  (       nn.Module						):
		UpperCamelCase						 =42
		UpperCamelCase						 =0.0
		UpperCamelCase						 =1
		UpperCamelCase						 =1
		UpperCamelCase						 =False
		UpperCamelCase						 =False
		UpperCamelCase						 =jnp.floataa
		def      _lowerCamelCase				(		self    )		->    Tuple:
					__lowercase       :   Optional[Any]        =    [
					    FlaxResnetBlockaD(
					        in_channels=self.in_channels							,		out_channels=self.in_channels							,		dropout_prob=self.dropout							,		dtype=self.dtype							,		)
					]

					__lowercase       :   Any        =    []

					for _ in range(self.num_layers    ):
								__lowercase       :   Optional[int]        =    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(lowercase_    )

								__lowercase       :   Tuple        =    FlaxResnetBlockaD(
								    in_channels=self.in_channels							,		out_channels=self.in_channels							,		dropout_prob=self.dropout							,		dtype=self.dtype							,		)
								resnets.append(lowercase_    )

					__lowercase       :   List[Any]        =    resnets
					__lowercase       :   List[str]        =    attentions
		def __call__(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_=True    )		->    int:
					__lowercase       :   Dict        =    self.resnets[0](lowercase_							,		lowercase_    )
					for attn, resnet in zip(self.attentions							,		self.resnets[1:]    ):
								__lowercase       :   Tuple        =    attn(lowercase_							,		lowercase_							,		deterministic=lowercase_    )
								__lowercase       :   List[str]        =    resnet(lowercase_							,		lowercase_							,		deterministic=lowercase_    )

					return hidden_states







import sys

from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

lowerCamelCase											=       'python tqdm regex requests packaging filelock numpy tokenizers'.split()
if sys.version_info < (3, 7):
 pkgs_to_check_at_runtime.append('dataclasses')
if sys.version_info < (3, 8):
 pkgs_to_check_at_runtime.append('importlib_metadata')

for pkg in pkgs_to_check_at_runtime:
 if pkg in deps:
  if pkg == "tokenizers":
   # must be loaded here, or else tqdm check may fail
   from .utils import is_tokenizers_available

   if not is_tokenizers_available():
    continue  # not required, check version only if installed

  require_version_core(deps[pkg])
 else:
  raise ValueError(F"""can't find {pkg} in {deps.keys()}, check dependency_versions_table.py""")




def 	a_      (    SCREAMING_SNAKE_CASE__					:      int		,			SCREAMING_SNAKE_CASE__					:      str=None     ):
       '''simple docstring'''





       require_version(deps[pkg]		,			SCREAMING_SNAKE_CASE__     )


