code: 
  
  import json
  import os
  import unittest
  
  from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES, FSMTTokenizer
  from transformers.testing_utils import slow
  from transformers.utils import cached_property
  
  from ...test_tokenization_common import TokenizerTesterMixin
  
  
  # using a different tiny model than the one used for default params defined in init to ensure proper testing
  A_		:			Dict										=   'stas/tiny-wmt19-en-ru'
  class 					A_   (      _a ,     unittest.TestCase						):
  
  
  				'''simple docstring'''
  
  
  
  
  
  
  				a__     							=       FSMTTokenizer
  				a__     							=       False
  
  
  				def      lowerCAmelCase_						(self  )     ->      List[Any]:
  								super().setUp()
  
  								# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
  								__UpperCAmelCase									=			[
  								    '''l''',
  								    '''o''',
  								    '''w''',
  								    '''e''',
  								    '''r''',
  								    '''s''',
  								    '''t''',
  								    '''i''',
  								    '''d''',
  								    '''n''',
  								    '''w</w>''',
  								    '''r</w>''',
  								    '''t</w>''',
  								    '''lo''',
  								    '''low''',
  								    '''er</w>''',
  								    '''low</w>''',
  								    '''lowest</w>''',
  								    '''newer</w>''',
  								    '''wider</w>''',
  								    '''<unk>''',
  								]
  								__UpperCAmelCase									=			dict(zip(lowercase__     ,      range(len(lowercase__  )  )  )  )
  								__UpperCAmelCase									=			['''l o 123''', '''lo w 1456''', '''e r</w> 1789''', '''''']
  
  								__UpperCAmelCase									=			['''en''', '''ru''']
  								__UpperCAmelCase									=			{
  								    '''langs''': self.langs,
  								    '''src_vocab_size''': 10,
  								    '''tgt_vocab_size''': 20,
  								}
  
  								__UpperCAmelCase									=			os.path.join(self.tmpdirname     ,      VOCAB_FILES_NAMES['''src_vocab_file''']  )
  								__UpperCAmelCase									=			os.path.join(self.tmpdirname     ,      VOCAB_FILES_NAMES['''tgt_vocab_file''']  )
  								__UpperCAmelCase									=			os.path.join(self.tmpdirname     ,      '''tokenizer_config.json'''  )
  								__UpperCAmelCase									=			os.path.join(self.tmpdirname     ,      VOCAB_FILES_NAMES['''merges_file''']  )
  								with open(self.src_vocab_file     ,      '''w'''  ) as fp:
  												fp.write(json.dumps(lowercase__  )  )
  								with open(self.tgt_vocab_file     ,      '''w'''  ) as fp:
  												fp.write(json.dumps(lowercase__  )  )
  								with open(self.merges_file     ,      '''w'''  ) as fp:
  												fp.write('''\n'''.join(lowercase__  )  )
  								with open(lowercase__     ,      '''w'''  ) as fp:
  												fp.write(json.dumps(lowercase__  )  )
  
  
  				@cached_property
  				def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
  								return FSMTTokenizer.from_pretrained('''facebook/wmt19-ru-en'''  )
  
  
  				@cached_property
  				def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
  								return FSMTTokenizer.from_pretrained('''facebook/wmt19-en-ru'''  )
  
  
  				def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
  								__UpperCAmelCase									=			FSMTTokenizer.from_pretrained(lowercase__  )
  								self.assertListEqual([tokenizer.src_lang, tokenizer.tgt_lang]     ,      ['''en''', '''ru''']  )
  								self.assertEqual(tokenizer.src_vocab_size     ,      21  )
  								self.assertEqual(tokenizer.tgt_vocab_size     ,      21  )
  
  
  				def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
  								__UpperCAmelCase									=			FSMTTokenizer(self.langs     ,      self.src_vocab_file     ,      self.tgt_vocab_file     ,      self.merges_file  )
  
  								__UpperCAmelCase									=			'''lower'''
  								__UpperCAmelCase									=			['''low''', '''er</w>''']
  								__UpperCAmelCase									=			tokenizer.tokenize(lowercase__  )
  								self.assertListEqual(lowercase__     ,      lowercase__  )
  
  								__UpperCAmelCase									=			tokens + ['''<unk>''']
  								__UpperCAmelCase									=			[14, 15, 20]
  								self.assertListEqual(tokenizer.convert_tokens_to_ids(lowercase__  )     ,      lowercase__  )
  
  
  				@slow
  				def      lowerCAmelCase_						(self  )     ->      Any:
  								__UpperCAmelCase									=			self.tokenizer_ru_en
  
  								__UpperCAmelCase									=			tokenizer.encode('''sequence builders'''     ,      add_special_tokens=lowercase__  )
  								__UpperCAmelCase									=			tokenizer.encode('''multi-sequence build'''     ,      add_special_tokens=lowercase__  )
  
  								__UpperCAmelCase									=			tokenizer.build_inputs_with_special_tokens(lowercase__  )
  								__UpperCAmelCase									=			tokenizer.build_inputs_with_special_tokens(lowercase__     ,      lowercase__  )
  
  								assert encoded_sentence == text + [2]
  								assert encoded_pair == text + [2] + text_a + [2]
  
  
  				@slow
  				def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
  								__UpperCAmelCase									=			self.tokenizer_en_ru
  								__UpperCAmelCase									=			self.tokenizer_ru_en
  
  								__UpperCAmelCase									=			[
  								    [
  								        '''Here\'s a little song I wrote. Don\'t worry, be happy.''',
  								        [2_470, 39, 11, 2_349, 7_222, 70, 5_979, 7, 8_450, 1_050, 13_160, 5, 26, 6_445, 7, 2],
  								    ],
  								    ['''This is it. No more. I\'m done!''', [132, 21, 37, 7, 1_434, 86, 7, 70, 6_476, 1_305, 427, 2]],
  								]
  
  								# if data needs to be recreated or added, run:
  								# import torch
  								# model = torch.hub.load("pytorch/fairseq", "transformer.wmt19.en-ru", checkpoint_file="model4.pt", tokenizer="moses", bpe="fastbpe")
  								# for src_text, _ in targets: print(f"""[\n"{src_text}",\n {model.encode(src_text).tolist()}\n],""")
  
  								for src_text, tgt_input_ids in targets:
  												__UpperCAmelCase									=			tokenizer_enc.encode(lowercase__     ,      return_tensors=lowercase__  )
  												self.assertListEqual(lowercase__     ,      lowercase__  )
  
  												# and decode backward, using the reversed languages model
  												__UpperCAmelCase									=			tokenizer_dec.decode(lowercase__     ,      skip_special_tokens=lowercase__  )
  												self.assertEqual(lowercase__     ,      lowercase__  )
  
  
  				@slow
  				def      lowerCAmelCase_						(self  )     ->      Tuple:
  								__UpperCAmelCase									=			FSMTTokenizer.from_pretrained('''facebook/wmt19-ru-en'''     ,      do_lower_case=lowercase__  )
  								__UpperCAmelCase									=			tokenizer.tokenize('''USA is United States of America'''  )
  								__UpperCAmelCase									=			['''us''', '''a</w>''', '''is</w>''', '''un''', '''i''', '''ted</w>''', '''st''', '''ates</w>''', '''of</w>''', '''am''', '''er''', '''ica</w>''']
  								self.assertListEqual(lowercase__     ,      lowercase__  )
  
  
  				@unittest.skip('''FSMTConfig.__init__  requires non-optional args'''  )
  				def      lowerCAmelCase_						(self  )     ->      Any:
  								pass
  
  
  
  
  
  
  				@unittest.skip('''FSMTConfig.__init__  requires non-optional args'''  )
  				def      lowerCAmelCase_						(self  )     ->      int:
  								pass
  
  
  
  
  
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__  )
  
  
  
  
  