


'''simple docstring'''



import collections
import importlib.util
import os
import re
from pathlib import Path


A__						:   Dict				=    """src/transformers"""


# Matches is_xxx_available()
A__						:   Union[str, Any]				=    re.compile(R"""is\_([a-z_]*)_available()""")
# Catches a one-line _import_struct = {xxx}
A__						:   Tuple				=    re.compile(R"""^_import_structure\s+=\s+\{([^\}]+)\}""")
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
A__						:   str				=    re.compile(R"""\s+\"\S*\":\s+\[([^\]]*)\]""")
# Catches a line if not is_foo_available
A__						:   Union[str, Any]				=    re.compile(R"""^\s*if\s+not\s+is\_[a-z_]*\_available\(\)""")
# Catches a line _import_struct["bla"].append("foo")
A__						:   int				=    re.compile(R"""^\s*_import_structure\[\"\S*\"\]\.append\(\"(\S*)\"\)""")
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
A__						:   List[Any]				=    re.compile(R"""^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]""")
# Catches a line with an object between quotes and a comma:     "MyModel",
A__						:   Tuple				=    re.compile("""^\s+\"([^\"]+)\",""")
# Catches a line with objects between brackets only:    ["foo", "bar"],
A__						:   List[Any]				=    re.compile("""^\s+\[([^\]]+)\]""")
# Catches a line with from foo import bar, bla, boo
A__						:   Any				=    re.compile(R"""\s+from\s+\S*\s+import\s+([^\(\s].*)\n""")
# Catches a line with try:
A__						:   str				=    re.compile(R"""^\s*try:""")
# Catches a line with else:
A__						:   List[str]				=    re.compile(R"""^\s*else:""")



def      UpperCAmelCase__   (      UpperCAmelCase_       :							List[Any]				)	->						List[str]:
     if _re_test_backend.search(UpperCAmelCase_				) is None:
          return None
     __lowerCamelCase  :      str														=						[b[0] for b in _re_backend.findall(UpperCAmelCase_				)]
     backends.sort()
     return "_and_".join(UpperCAmelCase_				)



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Tuple				)	->						Union[str, Any]:
     with open(UpperCAmelCase_ ,       'r' ,       encoding='utf-8' ,       newline='\n'				) as f:
          __lowerCamelCase  :      Union[str, Any]														=						f.readlines()

     __lowerCamelCase  :      Any														=						0
     while line_index < len(UpperCAmelCase_				) and not lines[line_index].startswith('_import_structure = {'				):
          line_index += 1

     # If this is a traditional init, just return.
     if line_index >= len(UpperCAmelCase_				):
          return None

     # First grab the objects without a specific backend in _import_structure
     __lowerCamelCase  :      Dict														=						[]
     while not lines[line_index].startswith('if TYPE_CHECKING'				) and find_backend(lines[line_index]				) is None:
          __lowerCamelCase  :      Optional[Any]														=						lines[line_index]
          # If we have everything on a single line, let's deal with it.
          if _re_one_line_import_struct.search(UpperCAmelCase_				):
               __lowerCamelCase  :      Optional[int]														=						_re_one_line_import_struct.search(UpperCAmelCase_				).groups()[0]
               __lowerCamelCase  :      Union[str, Any]														=						re.findall('\[([^\]]+)\]' ,       UpperCAmelCase_				)
               for imp in imports:
                    objects.extend([obj[1:-1] for obj in imp.split(', '				)]				)
               line_index += 1
               continue
          __lowerCamelCase  :      Tuple														=						_re_import_struct_key_value.search(UpperCAmelCase_				)
          if single_line_import_search is not None:
               __lowerCamelCase  :      Dict														=						[obj[1:-1] for obj in single_line_import_search.groups()[0].split(', '				) if len(UpperCAmelCase_				) > 0]
               objects.extend(UpperCAmelCase_				)
          elif line.startswith(' ' * 8 + '"'				):
               objects.append(line[9:-3]				)
          line_index += 1

     __lowerCamelCase  :      Tuple														=						{'none': objects}
     # Let's continue with backend-specific objects in _import_structure
     while not lines[line_index].startswith('if TYPE_CHECKING'				):
          # If the line is an if not is_backend_available, we grab all objects associated.
          __lowerCamelCase  :      Dict														=						find_backend(lines[line_index]				)
          # Check if the backend declaration is inside a try block:
          if _re_try.search(lines[line_index - 1]				) is None:
               __lowerCamelCase  :      int														=						None

          if backend is not None:
               line_index += 1

               # Scroll until we hit the else block of try-except-else
               while _re_else.search(lines[line_index]				) is None:
                    line_index += 1

               line_index += 1

               __lowerCamelCase  :      Tuple														=						[]
               # Until we unindent, add backend objects to the list
               while len(lines[line_index]				) <= 1 or lines[line_index].startswith(' ' * 4				):
                    __lowerCamelCase  :      int														=						lines[line_index]
                    if _re_import_struct_add_one.search(UpperCAmelCase_				) is not None:
                         objects.append(_re_import_struct_add_one.search(UpperCAmelCase_				).groups()[0]				)
                    elif _re_import_struct_add_many.search(UpperCAmelCase_				) is not None:
                         __lowerCamelCase  :      int														=						_re_import_struct_add_many.search(UpperCAmelCase_				).groups()[0].split(', '				)
                         __lowerCamelCase  :      Union[str, Any]														=						[obj[1:-1] for obj in imports if len(UpperCAmelCase_				) > 0]
                         objects.extend(UpperCAmelCase_				)
                    elif _re_between_brackets.search(UpperCAmelCase_				) is not None:
                         __lowerCamelCase  :      Dict														=						_re_between_brackets.search(UpperCAmelCase_				).groups()[0].split(', '				)
                         __lowerCamelCase  :      Union[str, Any]														=						[obj[1:-1] for obj in imports if len(UpperCAmelCase_				) > 0]
                         objects.extend(UpperCAmelCase_				)
                    elif _re_quote_object.search(UpperCAmelCase_				) is not None:
                         objects.append(_re_quote_object.search(UpperCAmelCase_				).groups()[0]				)
                    elif line.startswith(' ' * 8 + '"'				):
                         objects.append(line[9:-3]				)
                    elif line.startswith(' ' * 12 + '"'				):
                         objects.append(line[13:-3]				)
                    line_index += 1

               __lowerCamelCase  :      Optional[Any]														=						objects
          else:
               line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
     __lowerCamelCase  :      Union[str, Any]														=						[]
     while (
         line_index < len(UpperCAmelCase_				)
         and find_backend(lines[line_index]				) is None
         and not lines[line_index].startswith('else'				)
     ):
          __lowerCamelCase  :      List[str]														=						lines[line_index]
          __lowerCamelCase  :      Dict														=						_re_import.search(UpperCAmelCase_				)
          if single_line_import_search is not None:
               objects.extend(single_line_import_search.groups()[0].split(', '				)				)
          elif line.startswith(' ' * 8				):
               objects.append(line[8:-2]				)
          line_index += 1

     __lowerCamelCase  :      str														=						{'none': objects}
     # Let's continue with backend-specific objects
     while line_index < len(UpperCAmelCase_				):
          # If the line is an if is_backend_available, we grab all objects associated.
          __lowerCamelCase  :      Tuple														=						find_backend(lines[line_index]				)
          # Check if the backend declaration is inside a try block:
          if _re_try.search(lines[line_index - 1]				) is None:
               __lowerCamelCase  :      Optional[Any]														=						None

          if backend is not None:
               line_index += 1

               # Scroll until we hit the else block of try-except-else
               while _re_else.search(lines[line_index]				) is None:
                    line_index += 1

               line_index += 1

               __lowerCamelCase  :      Dict														=						[]
               # Until we unindent, add backend objects to the list
               while len(lines[line_index]				) <= 1 or lines[line_index].startswith(' ' * 8				):
                    __lowerCamelCase  :      Optional[Any]														=						lines[line_index]
                    __lowerCamelCase  :      Optional[Any]														=						_re_import.search(UpperCAmelCase_				)
                    if single_line_import_search is not None:
                         objects.extend(single_line_import_search.groups()[0].split(', '				)				)
                    elif line.startswith(' ' * 12				):
                         objects.append(line[12:-2]				)
                    line_index += 1

               __lowerCamelCase  :      Any														=						objects
          else:
               line_index += 1

     return import_dict_objects, type_hint_objects



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Optional[Any] ,       UpperCAmelCase_       :							List[Any]				)	->						List[Any]:

     def find_duplicates(UpperCAmelCase_       :							List[str]				):
          return [k for k, v in collections.Counter(UpperCAmelCase_				).items() if v > 1]

     if list(import_dict_objects.keys()				) != list(type_hint_objects.keys()				):
          return ["Both sides of the init do not have the same backends!"]

     __lowerCamelCase  :      Dict														=						[]
     for key in import_dict_objects.keys():
          __lowerCamelCase  :      Any														=						find_duplicates(import_dict_objects[key]				)
          if duplicate_imports:
               errors.append(F'Duplicate _import_structure definitions for: {duplicate_imports}'				)
          __lowerCamelCase  :      Optional[Any]														=						find_duplicates(type_hint_objects[key]				)
          if duplicate_type_hints:
               errors.append(F'Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}'				)

          if sorted(set(import_dict_objects[key]				)				) != sorted(set(type_hint_objects[key]				)				):
               __lowerCamelCase  :      Tuple														=						'base imports' if key == 'none' else F'{key} backend'
               errors.append(F'Differences for {name}:'				)
               for a in type_hint_objects[key]:
                    if a not in import_dict_objects[key]:
                         errors.append(F'  {a} in TYPE_HINT but not in _import_structure.'				)
               for a in import_dict_objects[key]:
                    if a not in type_hint_objects[key]:
                         errors.append(F'  {a} in _import_structure but not in TYPE_HINT.'				)
     return errors



def      UpperCAmelCase__   (      )	->						str:
     __lowerCamelCase  :      str														=						[]
     for root, _, files in os.walk(UpperCAmelCase_				):
          if "__init__.py" in files:
               __lowerCamelCase  :      Dict														=						os.path.join(UpperCAmelCase_ ,       '__init__.py'				)
               __lowerCamelCase  :      int														=						parse_init(UpperCAmelCase_				)
               if objects is not None:
                    __lowerCamelCase  :      Optional[Any]														=						analyze_results(*UpperCAmelCase_				)
                    if len(UpperCAmelCase_				) > 0:
                         __lowerCamelCase  :      Optional[int]														=						F'Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'
                         failures.append('\n'.join(UpperCAmelCase_				)				)
     if len(UpperCAmelCase_				) > 0:
          raise ValueError('\n\n'.join(UpperCAmelCase_				)				)



def      UpperCAmelCase__   (      )	->						Union[str, Any]:
     __lowerCamelCase  :      List[Any]														=						[]
     for path, directories, files in os.walk(UpperCAmelCase_				):
          for folder in directories:
               # Ignore private modules
               if folder.startswith('_'				):
                    directories.remove(UpperCAmelCase_				)
                    continue
               # Ignore leftovers from branches (empty folders apart from pycache)
               if len(list((Path(UpperCAmelCase_				) / folder).glob('*.py'				)				)				) == 0:
                    continue
               __lowerCamelCase  :      List[str]														=						str((Path(UpperCAmelCase_				) / folder).relative_to(UpperCAmelCase_				)				)
               __lowerCamelCase  :      int														=						short_path.replace(os.path.sep ,       '.'				)
               submodules.append(UpperCAmelCase_				)
          for fname in files:
               if fname == "__init__.py":
                    continue
               __lowerCamelCase  :      int														=						str((Path(UpperCAmelCase_				) / fname).relative_to(UpperCAmelCase_				)				)
               __lowerCamelCase  :      int														=						short_path.replace('.py' ,       ''				).replace(os.path.sep ,       '.'				)
               if len(submodule.split('.'				)				) == 1:
                    submodules.append(UpperCAmelCase_				)
     return submodules


A__						:   Optional[Any]				=    [
    """convert_pytorch_checkpoint_to_tf2""",
    """modeling_flax_pytorch_utils""",
]



def      UpperCAmelCase__   (      )	->						List[Any]:
     # This is to make sure the transformers module imported is the one in the repo.
     __lowerCamelCase  :      List[str]														=						importlib.util.spec_from_file_location(
         'transformers' ,       os.path.join(UpperCAmelCase_ ,       '__init__.py'				) ,       submodule_search_locations=[PATH_TO_TRANSFORMERS] ,       )
     __lowerCamelCase  :      Dict														=						spec.loader.load_module()

     __lowerCamelCase  :      str														=						[
         module
         for module in get_transformers_submodules()
         if module not in IGNORE_SUBMODULES and module not in transformers._import_structure.keys()
     ]
     if len(UpperCAmelCase_				) > 0:
          __lowerCamelCase  :      List[Any]														=						'\n'.join(F'- {module}' for module in module_not_registered				)
          raise ValueError(
              'The following submodules are not properly registered in the main init of Transformers:\n'
              F'{list_of_modules}\n'
              'Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value.'				)


if __name__ == "__main__":
      check_all_inits()
      check_submodules()










'''simple docstring'''



import flax.linen as nn
import jax
import jax.numpy as jnp





class   UpperCAmelCase_    (nn.Module    ):





    """simple docstring"""






    lowerCamelCase    :   int
    lowerCamelCase    :   jnp.dtype					=     jnp.floataa







    def  lowercase_							(				self      )							->	Union[str, Any]:
         __lowerCamelCase  :      Optional[Any]														=						nn.Conv(
             self.out_channels     ,	kernel_size=(3, 3)     ,	strides=(1, 1)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)







    def __call__(				self     ,	SCREAMING_SNAKE_CASE_      )							->	Any:
         __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase  :      Optional[int]														=						hidden_states.shape
         __lowerCamelCase  :      Dict														=						jax.image.resize(
             SCREAMING_SNAKE_CASE_     ,	shape=(batch, height * 2, width * 2, channels)     ,	method='nearest'     ,	)
         __lowerCamelCase  :      Optional[Any]														=						self.conv(SCREAMING_SNAKE_CASE_      )
         return hidden_states







class   UpperCAmelCase_    (nn.Module    ):





    """simple docstring"""






    lowerCamelCase    :   int
    lowerCamelCase    :   jnp.dtype					=     jnp.floataa







    def  lowercase_							(				self      )							->	List[str]:
         __lowerCamelCase  :      str														=						nn.Conv(
             self.out_channels     ,	kernel_size=(3, 3)     ,	strides=(2, 2)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)







    def __call__(				self     ,	SCREAMING_SNAKE_CASE_      )							->	List[str]:
         # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
         # hidden_states = jnp.pad(hidden_states, pad_width=pad)
         __lowerCamelCase  :      str														=						self.conv(SCREAMING_SNAKE_CASE_      )
         return hidden_states







class   UpperCAmelCase_    (nn.Module    ):





    """simple docstring"""






    lowerCamelCase    :   int
    lowerCamelCase    :   int					=     None
    lowerCamelCase    :   float					=     0.0
    lowerCamelCase    :   bool					=     None
    lowerCamelCase    :   jnp.dtype					=     jnp.floataa







    def  lowercase_							(				self      )							->	Optional[int]:
         __lowerCamelCase  :      Optional[Any]														=						self.in_channels if self.out_channels is None else self.out_channels

         __lowerCamelCase  :      Optional[Any]														=						nn.GroupNorm(num_groups=32     ,	epsilon=1E-5      )
         __lowerCamelCase  :      Tuple														=						nn.Conv(
             SCREAMING_SNAKE_CASE_     ,	kernel_size=(3, 3)     ,	strides=(1, 1)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)

         __lowerCamelCase  :      List[str]														=						nn.Dense(SCREAMING_SNAKE_CASE_     ,	dtype=self.dtype      )

         __lowerCamelCase  :      Dict														=						nn.GroupNorm(num_groups=32     ,	epsilon=1E-5      )
         __lowerCamelCase  :      int														=						nn.Dropout(self.dropout_prob      )
         __lowerCamelCase  :      Union[str, Any]														=						nn.Conv(
             SCREAMING_SNAKE_CASE_     ,	kernel_size=(3, 3)     ,	strides=(1, 1)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)

         __lowerCamelCase  :      Optional[int]														=						self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

         __lowerCamelCase  :      List[Any]														=						None
         if use_nin_shortcut:
              __lowerCamelCase  :      Any														=						nn.Conv(
                  SCREAMING_SNAKE_CASE_     ,	kernel_size=(1, 1)     ,	strides=(1, 1)     ,	padding='VALID'     ,	dtype=self.dtype     ,	)







    def __call__(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=True      )							->	Tuple:
         __lowerCamelCase  :      List[Any]														=						hidden_states
         __lowerCamelCase  :      str														=						self.norma(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Union[str, Any]														=						nn.swish(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Union[str, Any]														=						self.conva(SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      str														=						self.time_emb_proj(nn.swish(SCREAMING_SNAKE_CASE_      )      )
         __lowerCamelCase  :      List[str]														=						jnp.expand_dims(jnp.expand_dims(SCREAMING_SNAKE_CASE_     ,	1      )     ,	1      )
         __lowerCamelCase  :      Optional[int]														=						hidden_states + temb

         __lowerCamelCase  :      List[Any]														=						self.norma(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Any														=						nn.swish(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      List[Any]														=						self.dropout(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      List[Any]														=						self.conva(SCREAMING_SNAKE_CASE_      )

         if self.conv_shortcut is not None:
              __lowerCamelCase  :      List[str]														=						self.conv_shortcut(SCREAMING_SNAKE_CASE_      )

         return hidden_states + residual






