


"""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 _lowerCAmelCase    (				nn.Module     ):
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	float							  =  0.0
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	bool							  =  True
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	jnp.dtype							  =  jnp.floataa







   def 			lowerCamelCase		(     self						)     ->		List[str]:

       '''simple docstring'''




       snake_case   :				List[Any]				 =   []
       snake_case   :				List[str]				 =   []

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

           snake_case   :				List[str]				 =   FlaxResnetBlockaD(
               in_channels=UpperCamelCase__						,     out_channels=self.out_channels						,     dropout_prob=self.dropout						,     dtype=self.dtype						,     )
           resnets.append(UpperCamelCase__						)

           snake_case   :				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(UpperCamelCase__						)

       snake_case   :				Union[str, Any]				 =   resnets
       snake_case   :				Dict				 =   attentions

       if self.add_downsample:
           snake_case   :				Optional[Any]				 =   FlaxDownsampleaD(self.out_channels						,     dtype=self.dtype						)







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

       '''simple docstring'''




       snake_case   :				Optional[int]				 =   ()

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

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

       return hidden_states, output_states







class _lowerCAmelCase    (				nn.Module     ):
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	float							  =  0.0
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	bool							  =  True
   __UpperCAmelCase			:	jnp.dtype							  =  jnp.floataa







   def 			lowerCamelCase		(     self						)     ->		int:

       '''simple docstring'''




       snake_case   :				str				 =   []

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

           snake_case   :				List[str]				 =   FlaxResnetBlockaD(
               in_channels=UpperCamelCase__						,     out_channels=self.out_channels						,     dropout_prob=self.dropout						,     dtype=self.dtype						,     )
           resnets.append(UpperCamelCase__						)
       snake_case   :				Optional[Any]				 =   resnets

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







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

       '''simple docstring'''




       snake_case   :				str				 =   ()

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

       if self.add_downsample:
           snake_case   :				List[str]				 =   self.downsamplers_a(UpperCamelCase__						)
           output_states += (hidden_states,)

       return hidden_states, output_states







class _lowerCAmelCase    (				nn.Module     ):
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	float							  =  0.0
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	bool							  =  True
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	jnp.dtype							  =  jnp.floataa







   def 			lowerCamelCase		(     self						)     ->		Tuple:

       '''simple docstring'''




       snake_case   :				int				 =   []
       snake_case   :				int				 =   []

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

           snake_case   :				Tuple				 =   FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels						,     out_channels=self.out_channels						,     dropout_prob=self.dropout						,     dtype=self.dtype						,     )
           resnets.append(UpperCamelCase__						)

           snake_case   :				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(UpperCamelCase__						)

       snake_case   :				Dict				 =   resnets
       snake_case   :				Any				 =   attentions

       if self.add_upsample:
           snake_case   :				List[Any]				 =   FlaxUpsampleaD(self.out_channels						,     dtype=self.dtype						)







   def __call__(     self						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__=True						)     ->		Optional[Any]:

       '''simple docstring'''




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

           snake_case   :				Optional[int]				 =   resnet(UpperCamelCase__						,     UpperCamelCase__						,     deterministic=UpperCamelCase__						)
           snake_case   :				int				 =   attn(UpperCamelCase__						,     UpperCamelCase__						,     deterministic=UpperCamelCase__						)

       if self.add_upsample:
           snake_case   :				Any				 =   self.upsamplers_a(UpperCamelCase__						)

       return hidden_states







class _lowerCAmelCase    (				nn.Module     ):
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	float							  =  0.0
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	bool							  =  True
   __UpperCAmelCase			:	jnp.dtype							  =  jnp.floataa







   def 			lowerCamelCase		(     self						)     ->		str:

       '''simple docstring'''




       snake_case   :				Optional[int]				 =   []

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

           snake_case   :				Tuple				 =   FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels						,     out_channels=self.out_channels						,     dropout_prob=self.dropout						,     dtype=self.dtype						,     )
           resnets.append(UpperCamelCase__						)

       snake_case   :				Any				 =   resnets

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







   def __call__(     self						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__=True						)     ->		List[Any]:

       '''simple docstring'''




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

           snake_case   :				Optional[Any]				 =   resnet(UpperCamelCase__						,     UpperCamelCase__						,     deterministic=UpperCamelCase__						)

       if self.add_upsample:
           snake_case   :				List[str]				 =   self.upsamplers_a(UpperCamelCase__						)

       return hidden_states







class _lowerCAmelCase    (				nn.Module     ):
   __UpperCAmelCase			:	int
   __UpperCAmelCase			:	float							  =  0.0
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	int							  =  1
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	bool							  =  False
   __UpperCAmelCase			:	jnp.dtype							  =  jnp.floataa







   def 			lowerCamelCase		(     self						)     ->		Optional[int]:

       '''simple docstring'''




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

       snake_case   :				Dict				 =   []

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

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

       snake_case   :				Any				 =   resnets
       snake_case   :				Any				 =   attentions







   def __call__(     self						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__=True						)     ->		Optional[int]:

       '''simple docstring'''




       snake_case   :				List[str]				 =   self.resnets[0](UpperCamelCase__						,     UpperCamelCase__						)
       for attn, resnet in zip(self.attentions						,     self.resnets[1:]						):
           snake_case   :				List[Any]				 =   attn(UpperCamelCase__						,     UpperCamelCase__						,     deterministic=UpperCamelCase__						)
           snake_case   :				Tuple				 =   resnet(UpperCamelCase__						,     UpperCamelCase__						,     deterministic=UpperCamelCase__						)

       return hidden_states









"""simple docstring"""



import doctest
import logging
import os
import unittest
from pathlib import Path
from typing import List, Union

import transformers
from transformers.testing_utils import require_tf, require_torch, slow


__snake_case             =   logging.getLogger()







@unittest.skip('''Temporarily disable the doc tests.'''     )
@require_torch
@require_tf
@slow
class _lowerCAmelCase    (				unittest.TestCase     ):







   def 			lowerCamelCase		(     self						,     UpperCamelCase__						,     UpperCamelCase__ = None						,     UpperCamelCase__ = None						,     UpperCamelCase__ = None						,     UpperCamelCase__ = True						,     )     ->		Tuple:

       '''simple docstring'''




       snake_case   :				List[Any]				 =   [file for file in os.listdir(UpperCamelCase__						) if os.path.isfile(os.path.join(UpperCamelCase__						,     UpperCamelCase__						)						)]

       if identifier is not None:
           snake_case   :				Optional[Any]				 =   [file for file in files if identifier in file]

       if n_identifier is not None:
           if isinstance(UpperCamelCase__						,     UpperCamelCase__						):
               for n_ in n_identifier:
                   snake_case   :				str				 =   [file for file in files if n_ not in file]
           else:
               snake_case   :				str				 =   [file for file in files if n_identifier not in file]

       snake_case   :				Tuple				 =   ignore_files or []
       ignore_files.append("__init__.py"						)
       snake_case   :				Optional[Any]				 =   [file for file in files if file not in ignore_files]

       for file in files:
           # Open all files
           print("Testing"						,     UpperCamelCase__						)

           if only_modules:
               snake_case   :				str				 =   file.split("."						)[0]
               try:
                   snake_case   :				Optional[int]				 =   getattr(UpperCamelCase__						,     UpperCamelCase__						)
                   snake_case   :				str				 =   doctest.DocTestSuite(UpperCamelCase__						)
                   snake_case   :				Optional[Any]				 =   unittest.TextTestRunner().run(UpperCamelCase__						)
                   self.assertIs(len(result.failures						)						,     0						)
               except AttributeError:
                   logger.info(F'{module_identifier} is not a module.'						)
           else:
               snake_case   :				Tuple				 =   doctest.testfile(str(".." / directory / file						)						,     optionflags=doctest.ELLIPSIS						)
               self.assertIs(result.failed						,     0						)







   def 			lowerCamelCase		(     self						)     ->		str:

       '''simple docstring'''




       snake_case   :				Tuple				 =   Path("src/transformers"						)
       snake_case   :				List[Any]				 =   "modeling"
       snake_case   :				Dict				 =   [
           "modeling_ctrl.py",
           "modeling_tf_ctrl.py",
       ]
       self.analyze_directory(UpperCamelCase__						,     identifier=UpperCamelCase__						,     ignore_files=UpperCamelCase__						)







   def 			lowerCamelCase		(     self						)     ->		Dict:

       '''simple docstring'''




       snake_case   :				Optional[Any]				 =   Path("src/transformers"						)
       snake_case   :				Optional[Any]				 =   "tokenization"
       self.analyze_directory(UpperCamelCase__						,     identifier=UpperCamelCase__						)







   def 			lowerCamelCase		(     self						)     ->		List[str]:

       '''simple docstring'''




       snake_case   :				Union[str, Any]				 =   Path("src/transformers"						)
       snake_case   :				Optional[Any]				 =   "configuration"
       self.analyze_directory(UpperCamelCase__						,     identifier=UpperCamelCase__						)







   def 			lowerCamelCase		(     self						)     ->		Optional[Any]:

       '''simple docstring'''




       snake_case   :				List[str]				 =   Path("src/transformers"						)
       snake_case   :				List[str]				 =   ["configuration", "modeling", "tokenization"]
       self.analyze_directory(UpperCamelCase__						,     n_identifier=UpperCamelCase__						)







   def 			lowerCamelCase		(     self						)     ->		List[Any]:

       '''simple docstring'''




       snake_case   :				Any				 =   Path("docs/source"						)
       snake_case   :				Tuple				 =   ["favicon.ico"]
       self.analyze_directory(UpperCamelCase__						,     ignore_files=UpperCamelCase__						,     only_modules=UpperCamelCase__						)





