code: 
  
  # DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
  
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  
  import flax
  import jax.numpy as jnp
  from jax import random
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left
  @flax.struct.dataclass
  class 					A_   :
  
  
  				'''simple docstring'''
  
  
  
  
  
  
  				a__     							=       None
  				a__     							=       None
  				a__     							=       None
  
  
  				@classmethod
  				def      lowerCAmelCase_						(cls  )     ->      Optional[Any]:
  								return cls()
  
  
  
  
  
  
  @dataclass
  class 					A_   (      lowerCAmelCase__						):
  
  
  				'''simple docstring'''
  
  
  
  
  
  
  				a__     							=       42
  				a__     							=       42
  				a__     							=       None
  
  
  
  
  
  
  class 					A_   (      lowerCAmelCase__ ,     lowerCAmelCase__						):
  
  
  				'''simple docstring'''
  
  
  
  
  
  
  				@property
  				def      lowerCAmelCase_						(self  )     ->      Dict:
  								return True
  
  
  				@register_to_config
  				def __init__(self     ,      lowercase__ = 2_000     ,      lowercase__ = 0.15     ,      lowercase__ = 0.01     ,      lowercase__ = 1_348.0     ,      lowercase__ = 1E-5     ,      lowercase__ = 1     ,      )     ->      Optional[Any]:
  								pass
  
  
  				def      lowerCAmelCase_						(self  )     ->      Optional[int]:
  								__UpperCAmelCase									=			ScoreSdeVeSchedulerState.create()
  								return self.set_sigmas(
  								    lowerCAmelCase_     ,      self.config.num_train_timesteps     ,      self.config.sigma_min     ,      self.config.sigma_max     ,      self.config.sampling_eps     ,      )
  
  
  				def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__ = ()     ,      lowercase__ = None  )     ->      ScoreSdeVeSchedulerState:
  								__UpperCAmelCase									=			sampling_eps if sampling_eps is not None else self.config.sampling_eps
  
  								__UpperCAmelCase									=			jnp.linspace(1     ,      lowerCAmelCase_     ,      lowerCAmelCase_  )
  								return state.replace(timesteps=lowerCAmelCase_  )
  
  
  				def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__ = None     ,      lowercase__ = None     ,      lowercase__ = None     ,      )     ->      ScoreSdeVeSchedulerState:
  								__UpperCAmelCase									=			sigma_min if sigma_min is not None else self.config.sigma_min
  								__UpperCAmelCase									=			sigma_max if sigma_max is not None else self.config.sigma_max
  								__UpperCAmelCase									=			sampling_eps if sampling_eps is not None else self.config.sampling_eps
  								if state.timesteps is None:
  												__UpperCAmelCase									=			self.set_timesteps(lowerCAmelCase_     ,      lowerCAmelCase_     ,      lowerCAmelCase_  )
  
  								__UpperCAmelCase									=			jnp.exp(jnp.linspace(jnp.log(lowerCAmelCase_  )     ,      jnp.log(lowerCAmelCase_  )     ,      lowerCAmelCase_  )  )
  								__UpperCAmelCase									=			jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps]  )
  
  								return state.replace(discrete_sigmas=lowerCAmelCase_     ,      sigmas=lowerCAmelCase_  )
  
  
  				def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__  )     ->      List[Any]:
  								return jnp.where(timesteps == 0     ,      jnp.zeros_like(lowerCAmelCase_  )     ,      state.discrete_sigmas[timesteps - 1]  )
  
  
  				def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__ = True     ,      )     ->      Union[FlaxSdeVeOutput, Tuple]:
  
  
  
  
  
  
  								if state.timesteps is None:
  												raise ValueError(
  												    '''`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'''  )
  
  								__UpperCAmelCase									=			timestep * jnp.ones(
  								    sample.shape[0]     ,      )
  								__UpperCAmelCase									=			(timestep * (len(state.timesteps  ) - 1)).long()
  
  								__UpperCAmelCase									=			state.discrete_sigmas[timesteps]
  								__UpperCAmelCase									=			self.get_adjacent_sigma(lowerCAmelCase_     ,      lowerCAmelCase_     ,      lowerCAmelCase_  )
  								__UpperCAmelCase									=			jnp.zeros_like(lowerCAmelCase_  )
  								__UpperCAmelCase									=			(sigma**2 - adjacent_sigma**2) ** 0.5
  
  								# equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
  								# also equation 47 shows the analog from SDE models to ancestral sampling methods
  								__UpperCAmelCase									=			diffusion.flatten()
  								__UpperCAmelCase									=			broadcast_to_shape_from_left(lowerCAmelCase_     ,      sample.shape  )
  								__UpperCAmelCase									=			drift - diffusion**2 * model_output
  
  								#  equation 6: sample noise for the diffusion term of
  								__UpperCAmelCase									=			random.split(lowerCAmelCase_     ,      num=1  )
  								__UpperCAmelCase									=			random.normal(key=lowerCAmelCase_     ,      shape=sample.shape  )
  								__UpperCAmelCase									=			sample - drift  # subtract because `dt` is a small negative timestep
  								# TODO is the variable diffusion the correct scaling term for the noise?
  								__UpperCAmelCase									=			prev_sample_mean + diffusion * noise  # add impact of diffusion field g
  
  								if not return_dict:
  												return (prev_sample, prev_sample_mean, state)
  
  								return FlaxSdeVeOutput(prev_sample=lowerCAmelCase_     ,      prev_sample_mean=lowerCAmelCase_     ,      state=lowerCAmelCase_  )
  
  
  				def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__ = True     ,      )     ->      Union[FlaxSdeVeOutput, Tuple]:
  
  
  
  
  
  
  								if state.timesteps is None:
  												raise ValueError(
  												    '''`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'''  )
  
  								# For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
  								# sample noise for correction
  								__UpperCAmelCase									=			random.split(lowerCAmelCase_     ,      num=1  )
  								__UpperCAmelCase									=			random.normal(key=lowerCAmelCase_     ,      shape=sample.shape  )
  
  								# compute step size from the model_output, the noise, and the snr
  								__UpperCAmelCase									=			jnp.linalg.norm(lowerCAmelCase_  )
  								__UpperCAmelCase									=			jnp.linalg.norm(lowerCAmelCase_  )
  								__UpperCAmelCase									=			(self.config.snr * noise_norm / grad_norm) ** 2 * 2
  								__UpperCAmelCase									=			step_size * jnp.ones(sample.shape[0]  )
  
  								# compute corrected sample: model_output term and noise term
  								__UpperCAmelCase									=			step_size.flatten()
  								__UpperCAmelCase									=			broadcast_to_shape_from_left(lowerCAmelCase_     ,      sample.shape  )
  								__UpperCAmelCase									=			sample + step_size * model_output
  								__UpperCAmelCase									=			prev_sample_mean + ((step_size * 2) ** 0.5) * noise
  
  								if not return_dict:
  												return (prev_sample, state)
  
  								return FlaxSdeVeOutput(prev_sample=lowerCAmelCase_     ,      state=lowerCAmelCase_  )
  
  
  
  
  
  
  				def __len__(self  )     ->      Dict:
  								return self.config.num_train_timesteps
  
  
  
  
  
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__  )
  
  
  
  
  