code: 
  
  import itertools
  import os
  import subprocess
  from os.path import dirname
  
  from parameterized import parameterized
  
  from tests.trainer.test_trainer import TrainerIntegrationCommon  # noqa
  from transformers import is_torch_available
  from transformers.testing_utils import (
      TestCasePlus,
      execute_subprocess_async,
      get_gpu_count,
      get_tests_dir,
      require_deepspeed,
      require_torch_gpu,
      slow,
  )
  from transformers.trainer_utils import set_seed
  
  
  if is_torch_available():
         from tests.trainer.test_trainer import (  # noqa
             RegressionModelConfig,
             RegressionPreTrainedModel,
             get_regression_trainer,
         )
  
  
  set_seed(42)
  
  A_		:			List[Any]										=   get_tests_dir('fixtures')
  A_		:			Dict										=   os.path.join(dirname(get_tests_dir()))
  A_		:			Optional[Any]										=   dirname(os.path.abspath(__file__))
  
  # default torch.distributed port
  A_		:			Optional[int]										=   '10999'
  
  A_		:			List[str]										=   't5-small'
  
  # *** Working Models ***
  A_		:			int										=   'hf-internal-testing/tiny-albert'
  A_		:			Union[str, Any]										=   'sshleifer/bart-tiny-random'
  A_		:			str										=   'hf-internal-testing/tiny-bert'
  A_		:			str										=   'hf-internal-testing/tiny-random-bigbird_pegasus'
  A_		:			Dict										=   'hf-internal-testing/tiny-random-big_bird'
  A_		:			str										=   'hf-internal-testing/tiny-random-blenderbot'
  A_		:			List[Any]										=   'bigscience/bigscience-small-testing'
  A_		:			Optional[int]										=   'hf-internal-testing/tiny-random-deberta'
  A_		:			Dict										=   'hf-internal-testing/tiny-random-deberta-v2'
  A_		:			Optional[Any]										=   'sshleifer/tiny-distilbert-base-cased'
  A_		:			Any										=   'hf-internal-testing/tiny-electra'
  A_		:			Optional[int]										=   'hf-internal-testing/tiny-random-flaubert'
  A_		:			Optional[int]										=   'stas/tiny-wmt19-en-de'
  A_		:			Union[str, Any]										=   'hf-internal-testing/tiny-random-funnel'
  A_		:			List[str]										=   'sshleifer/tiny-gpt2'
  A_		:			Optional[int]										=   'hf-internal-testing/tiny-random-gptj'
  A_		:			Union[str, Any]										=   'hf-internal-testing/tiny-random-gpt_neo'
  A_		:			int										=   'hf-internal-testing/tiny-layoutlm'
  A_		:			Tuple										=   'hf-internal-testing/tiny-random-led'
  A_		:			int										=   'hf-internal-testing/tiny-random-longformer'
  A_		:			List[Any]										=   'stas/tiny-m2m_100'  # hf tiny model is unsuitable
  A_		:			int										=   'sshleifer/tiny-marian-en-de'
  A_		:			List[str]										=   'sshleifer/tiny-mbart'
  A_		:			Tuple										=   'hf-internal-testing/tiny-random-mobilebert'
  A_		:			int										=   'hf-internal-testing/tiny-random-mpnet'
  A_		:			Optional[int]										=   'stas/pegasus-cnn_dailymail-tiny-random'
  A_		:			int										=   'hf-internal-testing/tiny-random-prophetnet'
  A_		:			Union[str, Any]										=   'sshleifer/tiny-distilroberta-base'
  A_		:			Any										=   'hf-internal-testing/tiny-random-squeezebert'
  A_		:			Optional[Any]										=   'patrickvonplaten/t5-tiny-random'
  A_		:			Optional[Any]										=   'hf-internal-testing/tiny-random-t5-v1.1'
  A_		:			Optional[int]										=   'hf-internal-testing/tiny-random-vit'
  A_		:			int										=   'hf-internal-testing/tiny-xlm-roberta'
  A_		:			Dict										=   'sshleifer/tiny-xlnet-base-cased'
  
  
  # *** To Fix ***
  
  
  # *** tiny model issues ***
  # missing model files:
  A_		:			Optional[Any]										=   'hf-internal-testing/tiny-random-mt5'
  A_		:			Tuple										=   'hf-internal-testing/tiny-random-camembert'
  A_		:			Optional[Any]										=   'hf-internal-testing/tiny-random-openai-gpt'
  
  # missing tokenizer files
  A_		:			str										=   'hf-internal-testing/tiny-random-convbert'
  A_		:			Tuple										=   'hf-internal-testing/tiny-random-layoutlmv2'
  A_		:			int										=   'hf-internal-testing/tiny-random-hubert'
  
  # issues with tokenizer
  A_		:			Tuple										=   'hf-internal-testing/tiny-random-ctrl'
  A_		:			List[str]										=   'hf-internal-testing/tiny-random-transfo-xl'  # same as ctrl
  
  # other issues with tiny models
  A_		:			Optional[Any]										=   'hf-internal-testing/tiny-random-ibert'  # multiple issues with either mlm/qa/clas
  A_		:			Optional[int]										=   'hf-internal-testing/tiny-random-reformer'  # multiple issues with either mlm/qa/clas
  
  # *** Lacking official examples to test with ***
  # or not working with examples
  A_		:			Any										=   'hf-internal-testing/tiny-random-dpr'
  # - "dpr"  examples/research_projects/rag-end2end-retriever/
  A_		:			Optional[Any]										=   'hf-internal-testing/tiny-random-rag'
  # - "rag" research_projects
  A_		:			Union[str, Any]										=   ''
  # - "luke" Entities classes - no plan to make such example
  A_		:			List[Any]										=   'hf-internal-testing/tiny-random-lxmert'
  # - "lxmert" doesn't work with run_qa.py
  A_		:			Optional[Any]										=   'hf-internal-testing/tiny-random-clip'
  # - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
  A_		:			List[str]										=   'hf-internal-testing/tiny-random-speech_to_text'
  # - "speech_to_text", nothing under pytorch examples
  
  
  # *** Reactive mode ***
  # models with low usage, unstable API, things about to change - do nothing about the following until someone runs into a problem
  A_		:			Optional[int]										=   'hf-internal-testing/tiny-random-tapas'
  
  
  def       __a       (   SCREAMING_SNAKE_CASE=False      )      ->     Optional[int]:
  
  
  
  
  
  
      '''simple docstring'''
      # 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
      # - it won't be able to handle that
      # 2. for now testing with just 2 gpus max (since some quality tests may give different
      # results with mode gpus because we use very little data)
      __UpperCAmelCase									=			min(2       ,   get_gpu_count()      ) if distributed else 1
      __UpperCAmelCase									=			os.environ.get('''DS_TEST_PORT'''       ,   SCREAMING_SNAKE_CASE      )
      return f'''deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'''.split()
  
  
  def       __a       (   )      ->     Union[str, Any]:
  
  
  
  
  
  
      '''simple docstring'''
      __UpperCAmelCase									=			f'''{FIXTURE_DIRECTORY}/tests_samples'''
      __UpperCAmelCase									=			f'''{data_dir_samples}/wmt_en_ro'''
      __UpperCAmelCase									=			f'''{data_dir_samples}/xsum'''
      __UpperCAmelCase									=			'''
          --do_train
          --max_train_samples 4
          --per_device_train_batch_size 2
          --num_train_epochs 1
          --fp16
          --report_to none
          --overwrite_output_dir
          '''.split()
  
      # try to cover as many models as possible once (it's enough to run on one task per model)
      # but need a tiny model for each
      #
      # should have "{model_type.upper()}_TINY" corresponding vars defined, e.g., T5_TINY, etc.
      __UpperCAmelCase									=			{
          '''trans''': [
              '''bart''',
              '''fsmt''',
              '''m2m_100''',
              '''marian''',
              '''mbart''',
              '''t5''',
              '''t5_v1''',
              # "mt5", missing model files
          ],
          '''sum''': [
              '''pegasus''',
          ],
          '''clm''': [
              '''big_bird''',
              '''bigbird_pegasus''',
              '''blenderbot''',
              '''bloom''',
              '''gpt2''',
              '''gpt_neo''',
              '''gptj''',
              '''xlm-roberta''',
              '''prophetnet''',
              # "camembert", missing model files
          ],
          '''mlm''': [
              '''albert''',
              '''deberta''',
              '''deberta-v2''',
              '''distilbert''',
              '''electra''',
              '''flaubert''',
              '''funnel''',
              '''layoutlm''',
              # "reformer", # multiple issues with either mlm/qa/clas
          ],
          '''qa''': [
              '''led''',
              '''longformer''',
              '''mobilebert''',
              '''mpnet''',
              '''roberta''',
              '''squeezebert''',
              # "convbert", # missing tokenizer files
              # "layoutlmv2", missing model files
          ],
          '''clas''': [
              '''bert''',
              '''xlnet''',
              # "hubert", # missing tokenizer files
              # "ibert", # multiple issues with either mlm/qa/clas
              # "transfo-xl", # tokenizer issues as ctrl
              # "ctrl", # tokenizer issues
              # "openai-gpt", missing model files
              # "tapas", multiple issues
          ],
          '''img_clas''': [
              '''vit''',
          ],
      }
  
      __UpperCAmelCase									=			f'''{ROOT_DIRECTORY}/examples/pytorch'''
  
      __UpperCAmelCase									=			{
          '''trans''': f'''
          {scripts_dir}/translation/run_translation.py
          --train_file {data_dir_wmt}/train.json
          --source_lang en
          --target_lang ro
          ''',
          '''sum''': f'''
          {scripts_dir}/summarization/run_summarization.py
          --train_file {data_dir_xsum}/sample.json
          --max_source_length 12
          --max_target_length 12
          --lang en
          ''',
          '''clm''': f'''
          {scripts_dir}/language-modeling/run_clm.py
          --train_file {FIXTURE_DIRECTORY}/sample_text.txt
          --block_size 8
          ''',
          '''mlm''': f'''
          {scripts_dir}/language-modeling/run_mlm.py
          --train_file {FIXTURE_DIRECTORY}/sample_text.txt
          ''',
          '''qa''': f'''
          {scripts_dir}/question-answering/run_qa.py
          --train_file {data_dir_samples}/SQUAD/sample.json
          ''',
          '''clas''': f'''
          {scripts_dir}/text-classification/run_glue.py
          --train_file {data_dir_samples}/MRPC/train.csv
          --max_seq_length 12
          --task_name MRPC
          ''',
          '''img_clas''': f'''
          {scripts_dir}/image-classification/run_image_classification.py
              --dataset_name hf-internal-testing/cats_vs_dogs_sample
              --remove_unused_columns False
              --max_steps 10
              --image_processor_name {DS_TESTS_DIRECTORY}/vit_feature_extractor.json
          ''',
      }
  
      __UpperCAmelCase									=			get_launcher(distributed=SCREAMING_SNAKE_CASE      )
  
      __UpperCAmelCase									=			{}
      for task, args in tasks.items():
          __UpperCAmelCase									=			args.split()
          for model in tasksamodels[task]:
              __UpperCAmelCase									=			globals()[f'''{model.upper().replace('-'       ,   '_'      )}_TINY''']
              __UpperCAmelCase									=			f'''--model_name_or_path {model_name}'''.split()
              __UpperCAmelCase									=			launcher + args + args_model + args_main
  
              # # generation special case
              # if task == "gen":
              #     launcher = f"deepspeed --num_nodes 1 --num_gpus 1".split()
              #     args_model += f"--model_type {model}".split()
              #     cmds[f"{task}_{model}"] = launcher + args + args_model
              # else:
  
      return cmds
  
  
  A_		:			int										=   make_task_cmds()
  
  A_		:			int										=   'zero2'
  A_		:			Union[str, Any]										=   'zero3'
  
  A_		:			Optional[int]										=   [ZEROa, ZEROa]
  
  
  def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     List[Any]:
  
  
  
  
  
  
      '''simple docstring'''
      # customize the test name generator function as we want both params to appear in the sub-test
      # name, as by default it shows only the first param
      __UpperCAmelCase									=			parameterized.to_safe_name('''_'''.join(str(SCREAMING_SNAKE_CASE      ) for x in param.args      )      )
      return f'''{func.__name__}_{param_based_name}'''
  
  
  # Cartesian-product of zero stages with models to test
  A_		:			Any										=   list(itertools.product(stages, task_cmds.keys()))
  @slow
  @require_deepspeed
  @require_torch_gpu
  class 					A_   (      _a						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__  )     ->      Tuple:
          # return a ready to run train cmd
          if task not in task_cmds:
              raise ValueError(F'''don\'t know of task {task}, have {task_cmds.keys()}'''  )
  
          __UpperCAmelCase									=			task_cmds[task]
          __UpperCAmelCase									=			F'''--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json'''.split()
  
          __UpperCAmelCase									=			self.get_auto_remove_tmp_dir()
          __UpperCAmelCase									=			F'''--output_dir {output_dir}'''.split()
  
          cmd += args_ds + args_out
  
          return cmd, output_dir
  
  
  
  
  
  
      @parameterized.expand(lowercase__     ,      name_func=lowercase__  )
      def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__  )     ->      Any:
          # testing the ability to do a run followed by recovery of full fp32 weights
  
          __UpperCAmelCase					,     __UpperCAmelCase									=			self.get_task_cmd(lowercase__     ,      lowercase__  )
  
          # 1. generate the checkpoint
          cmd += "--save_steps 1".split()
          # keep for quick debug
          # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] + cmd)); die
          execute_subprocess_async(lowercase__     ,      env=self.get_env()  )
  
          # 2. test that the fp32 weights get reconsolidated
          __UpperCAmelCase									=			F'''{output_dir}/checkpoint-1'''
          __UpperCAmelCase									=			F'''{chkpt_dir}/out.bin'''
          __UpperCAmelCase									=			F'''{chkpt_dir}/zero_to_fp32.py {chkpt_dir} {recovered_model_path}'''
          # keep for quick debug
          # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
          subprocess.check_call(lowercase__     ,      shell=lowercase__  )
          assert os.path.exists(lowercase__  ), F'''{recovered_model_path} was not found'''
  
          # possibly could also test that the resulting saved model is usable but given that we use
          # random models we won't know if it's any good
  
  
  
  
  
style_context: 
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE = 4_0_9_6      )      ->     Any:
  
  
  
  
  
  
      '''simple docstring'''
      __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			key.shape[-3:]
      __UpperCAmelCase									=			value.shape[-1]
      __UpperCAmelCase									=			min(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )
      __UpperCAmelCase									=			query / jnp.sqrt(SCREAMING_SNAKE_CASE      )
  
      @functools.partial(jax.checkpoint       ,   prevent_cse=SCREAMING_SNAKE_CASE      )
      def summarize_chunk(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      ):
          __UpperCAmelCase									=			jnp.einsum('''...qhd,...khd->...qhk'''       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   precision=SCREAMING_SNAKE_CASE      )
  
          __UpperCAmelCase									=			jnp.max(SCREAMING_SNAKE_CASE       ,   axis=-1       ,   keepdims=SCREAMING_SNAKE_CASE      )
          __UpperCAmelCase									=			jax.lax.stop_gradient(SCREAMING_SNAKE_CASE      )
          __UpperCAmelCase									=			jnp.exp(attn_weights - max_score      )
  
          __UpperCAmelCase									=			jnp.einsum('''...vhf,...qhv->...qhf'''       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   precision=SCREAMING_SNAKE_CASE      )
          __UpperCAmelCase									=			jnp.einsum('''...qhk->...qh'''       ,   SCREAMING_SNAKE_CASE      )
  
          return (exp_values, exp_weights.sum(axis=-1      ), max_score)
  
      def chunk_scanner(SCREAMING_SNAKE_CASE      ):
          # julienne key array
          __UpperCAmelCase									=			jax.lax.dynamic_slice(
              operand=SCREAMING_SNAKE_CASE       ,   start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]       ,   slice_sizes=list(key.shape[:-3]      ) + [key_chunk_size, num_heads, k_features]       ,   )
  
          # julienne value array
          __UpperCAmelCase									=			jax.lax.dynamic_slice(
              operand=SCREAMING_SNAKE_CASE       ,   start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]       ,   slice_sizes=list(value.shape[:-3]      ) + [key_chunk_size, num_heads, v_features]       ,   )
  
          return summarize_chunk(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )
  
      __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			jax.lax.map(f=SCREAMING_SNAKE_CASE       ,   xs=jnp.arange(0       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      )
  
      __UpperCAmelCase									=			jnp.max(SCREAMING_SNAKE_CASE       ,   axis=0       ,   keepdims=SCREAMING_SNAKE_CASE      )
      __UpperCAmelCase									=			jnp.exp(chunk_max - global_max      )
  
      chunk_values *= jnp.expand_dims(SCREAMING_SNAKE_CASE       ,   axis=-1      )
      chunk_weights *= max_diffs
  
      __UpperCAmelCase									=			chunk_values.sum(axis=0      )
      __UpperCAmelCase									=			jnp.expand_dims(SCREAMING_SNAKE_CASE       ,   -1      ).sum(axis=0      )
  
      return all_values / all_weights
  
  
  def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE=jax.lax.Precision.HIGHEST       ,   SCREAMING_SNAKE_CASE = 1_0_2_4       ,   SCREAMING_SNAKE_CASE = 4_0_9_6      )      ->     int:
  
  
  
  
  
  
      '''simple docstring'''
      __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			query.shape[-3:]
  
      def chunk_scanner(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      ):
          # julienne query array
          __UpperCAmelCase									=			jax.lax.dynamic_slice(
              operand=SCREAMING_SNAKE_CASE       ,   start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]       ,   slice_sizes=list(query.shape[:-3]      ) + [min(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      ), num_heads, q_features]       ,   )
  
          return (
              chunk_idx + query_chunk_size,  # unused ignore it
              _query_chunk_attention(
                  query=SCREAMING_SNAKE_CASE       ,   key=SCREAMING_SNAKE_CASE       ,   value=SCREAMING_SNAKE_CASE       ,   precision=SCREAMING_SNAKE_CASE       ,   key_chunk_size=SCREAMING_SNAKE_CASE      ),
          )
  
      __UpperCAmelCase					,     __UpperCAmelCase									=			jax.lax.scan(
          f=SCREAMING_SNAKE_CASE       ,   init=0       ,   xs=SCREAMING_SNAKE_CASE       ,   length=math.ceil(num_q / query_chunk_size      )      )
  
      return jnp.concatenate(SCREAMING_SNAKE_CASE       ,   axis=-3      )  # fuse the chunked result back
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       8
      a__     							=       64
      a__     							=       0.0
      a__     							=       False
      a__     							=       jnp.floataa
  
  
      def      lowerCAmelCase_						(self  )     ->      Tuple:
          __UpperCAmelCase									=			self.dim_head * self.heads
          __UpperCAmelCase									=			self.dim_head**-0.5
  
          # Weights were exported with old names {to_q, to_k, to_v, to_out}
          __UpperCAmelCase									=			nn.Dense(lowercase__     ,      use_bias=lowercase__     ,      dtype=self.dtype     ,      name='''to_q'''  )
          __UpperCAmelCase									=			nn.Dense(lowercase__     ,      use_bias=lowercase__     ,      dtype=self.dtype     ,      name='''to_k'''  )
          __UpperCAmelCase									=			nn.Dense(lowercase__     ,      use_bias=lowercase__     ,      dtype=self.dtype     ,      name='''to_v'''  )
  
          __UpperCAmelCase									=			nn.Dense(self.query_dim     ,      dtype=self.dtype     ,      name='''to_out_0'''  )
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
      def      lowerCAmelCase_						(self     ,      lowercase__  )     ->      Optional[int]:
          __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			tensor.shape
          __UpperCAmelCase									=			self.heads
          __UpperCAmelCase									=			tensor.reshape(lowercase__     ,      lowercase__     ,      lowercase__     ,      dim // head_size  )
          __UpperCAmelCase									=			jnp.transpose(lowercase__     ,      (0, 2, 1, 3)  )
          __UpperCAmelCase									=			tensor.reshape(batch_size * head_size     ,      lowercase__     ,      dim // head_size  )
          return tensor
  
  
      def      lowerCAmelCase_						(self     ,      lowercase__  )     ->      Any:
          __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			tensor.shape
          __UpperCAmelCase									=			self.heads
          __UpperCAmelCase									=			tensor.reshape(batch_size // head_size     ,      lowercase__     ,      lowercase__     ,      lowercase__  )
          __UpperCAmelCase									=			jnp.transpose(lowercase__     ,      (0, 2, 1, 3)  )
          __UpperCAmelCase									=			tensor.reshape(batch_size // head_size     ,      lowercase__     ,      dim * head_size  )
          return tensor
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__=None     ,      lowercase__=True  )     ->      Optional[Any]:
          __UpperCAmelCase									=			hidden_states if context is None else context
  
          __UpperCAmelCase									=			self.query(lowercase__  )
          __UpperCAmelCase									=			self.key(lowercase__  )
          __UpperCAmelCase									=			self.value(lowercase__  )
  
          __UpperCAmelCase									=			self.reshape_heads_to_batch_dim(lowercase__  )
          __UpperCAmelCase									=			self.reshape_heads_to_batch_dim(lowercase__  )
          __UpperCAmelCase									=			self.reshape_heads_to_batch_dim(lowercase__  )
  
          if self.use_memory_efficient_attention:
              __UpperCAmelCase									=			query_states.transpose(1     ,      0     ,      2  )
              __UpperCAmelCase									=			key_states.transpose(1     ,      0     ,      2  )
              __UpperCAmelCase									=			value_states.transpose(1     ,      0     ,      2  )
  
              # this if statement create a chunk size for each layer of the unet
              # the chunk size is equal to the query_length dimension of the deepest layer of the unet
  
              __UpperCAmelCase									=			query_states.shape[-3]
              if flatten_latent_dim % 64 == 0:
                  __UpperCAmelCase									=			int(flatten_latent_dim / 64  )
              elif flatten_latent_dim % 16 == 0:
                  __UpperCAmelCase									=			int(flatten_latent_dim / 16  )
              elif flatten_latent_dim % 4 == 0:
                  __UpperCAmelCase									=			int(flatten_latent_dim / 4  )
              else:
                  __UpperCAmelCase									=			int(lowercase__  )
  
              __UpperCAmelCase									=			jax_memory_efficient_attention(
                  lowercase__     ,      lowercase__     ,      lowercase__     ,      query_chunk_size=lowercase__     ,      key_chunk_size=4_096 * 4  )
  
              __UpperCAmelCase									=			hidden_states.transpose(1     ,      0     ,      2  )
          else:
              # compute attentions
              __UpperCAmelCase									=			jnp.einsum('''b i d, b j d->b i j'''     ,      lowercase__     ,      lowercase__  )
              __UpperCAmelCase									=			attention_scores * self.scale
              __UpperCAmelCase									=			nn.softmax(lowercase__     ,      axis=2  )
  
              # attend to values
              __UpperCAmelCase									=			jnp.einsum('''b i j, b j d -> b i d'''     ,      lowercase__     ,      lowercase__  )
  
          __UpperCAmelCase									=			self.reshape_batch_dim_to_heads(lowercase__  )
          __UpperCAmelCase									=			self.proj_attn(lowercase__  )
          return self.dropout_layer(lowercase__     ,      deterministic=lowercase__  )
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       42
      a__     							=       42
      a__     							=       0.0
      a__     							=       False
      a__     							=       jnp.floataa
      a__     							=       False
  
  
      def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
          # self attention (or cross_attention if only_cross_attention is True)
          __UpperCAmelCase									=			FlaxAttention(
              self.dim     ,      self.n_heads     ,      self.d_head     ,      self.dropout     ,      self.use_memory_efficient_attention     ,      dtype=self.dtype  )
          # cross attention
          __UpperCAmelCase									=			FlaxAttention(
              self.dim     ,      self.n_heads     ,      self.d_head     ,      self.dropout     ,      self.use_memory_efficient_attention     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			FlaxFeedForward(dim=self.dim     ,      dropout=self.dropout     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.LayerNorm(epsilon=1E-5     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.LayerNorm(epsilon=1E-5     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.LayerNorm(epsilon=1E-5     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__     ,      lowercase__=True  )     ->      Optional[int]:
          # self attention
          __UpperCAmelCase									=			hidden_states
          if self.only_cross_attention:
              __UpperCAmelCase									=			self.attna(self.norma(lowercase__  )     ,      lowercase__     ,      deterministic=lowercase__  )
          else:
              __UpperCAmelCase									=			self.attna(self.norma(lowercase__  )     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			hidden_states + residual
  
          # cross attention
          __UpperCAmelCase									=			hidden_states
          __UpperCAmelCase									=			self.attna(self.norma(lowercase__  )     ,      lowercase__     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			hidden_states + residual
  
          # feed forward
          __UpperCAmelCase									=			hidden_states
          __UpperCAmelCase									=			self.ff(self.norma(lowercase__  )     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			hidden_states + residual
  
          return self.dropout_layer(lowercase__     ,      deterministic=lowercase__  )
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       42
      a__     							=       42
      a__     							=       1
      a__     							=       0.0
      a__     							=       False
      a__     							=       False
      a__     							=       jnp.floataa
      a__     							=       False
  
  
      def      lowerCAmelCase_						(self  )     ->      int:
          __UpperCAmelCase									=			nn.GroupNorm(num_groups=32     ,      epsilon=1E-5  )
  
          __UpperCAmelCase									=			self.n_heads * self.d_head
          if self.use_linear_projection:
              __UpperCAmelCase									=			nn.Dense(lowercase__     ,      dtype=self.dtype  )
          else:
              __UpperCAmelCase									=			nn.Conv(
                  lowercase__     ,      kernel_size=(1, 1)     ,      strides=(1, 1)     ,      padding='''VALID'''     ,      dtype=self.dtype     ,      )
  
          __UpperCAmelCase									=			[
              FlaxBasicTransformerBlock(
                  lowercase__     ,      self.n_heads     ,      self.d_head     ,      dropout=self.dropout     ,      only_cross_attention=self.only_cross_attention     ,      dtype=self.dtype     ,      use_memory_efficient_attention=self.use_memory_efficient_attention     ,      )
              for _ in range(self.depth  )
          ]
  
          if self.use_linear_projection:
              __UpperCAmelCase									=			nn.Dense(lowercase__     ,      dtype=self.dtype  )
          else:
              __UpperCAmelCase									=			nn.Conv(
                  lowercase__     ,      kernel_size=(1, 1)     ,      strides=(1, 1)     ,      padding='''VALID'''     ,      dtype=self.dtype     ,      )
  
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__     ,      lowercase__=True  )     ->      Dict:
          __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			hidden_states.shape
          __UpperCAmelCase									=			hidden_states
          __UpperCAmelCase									=			self.norm(lowercase__  )
          if self.use_linear_projection:
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      height * width     ,      lowercase__  )
              __UpperCAmelCase									=			self.proj_in(lowercase__  )
          else:
              __UpperCAmelCase									=			self.proj_in(lowercase__  )
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      height * width     ,      lowercase__  )
  
          for transformer_block in self.transformer_blocks:
              __UpperCAmelCase									=			transformer_block(lowercase__     ,      lowercase__     ,      deterministic=lowercase__  )
  
          if self.use_linear_projection:
              __UpperCAmelCase									=			self.proj_out(lowercase__  )
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__  )
          else:
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__  )
              __UpperCAmelCase									=			self.proj_out(lowercase__  )
  
          __UpperCAmelCase									=			hidden_states + residual
          return self.dropout_layer(lowercase__     ,      deterministic=lowercase__  )
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       0.0
      a__     							=       jnp.floataa
  
  
      def      lowerCAmelCase_						(self  )     ->      int:
          # The second linear layer needs to be called
          # net_2 for now to match the index of the Sequential layer
          __UpperCAmelCase									=			FlaxGEGLU(self.dim     ,      self.dropout     ,      self.dtype  )
          __UpperCAmelCase									=			nn.Dense(self.dim     ,      dtype=self.dtype  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__=True  )     ->      Dict:
          __UpperCAmelCase									=			self.net_a(lowercase__     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			self.net_a(lowercase__  )
          return hidden_states
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       0.0
      a__     							=       jnp.floataa
  
  
      def      lowerCAmelCase_						(self  )     ->      Optional[int]:
          __UpperCAmelCase									=			self.dim * 4
          __UpperCAmelCase									=			nn.Dense(inner_dim * 2     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__=True  )     ->      Union[str, Any]:
          __UpperCAmelCase									=			self.proj(lowercase__  )
          __UpperCAmelCase					,     __UpperCAmelCase									=			jnp.split(lowercase__     ,      2     ,      axis=2  )
          return self.dropout_layer(hidden_linear * nn.gelu(lowercase__  )     ,      deterministic=lowercase__  )
  
  
  
  
  