class_docstr: A Spatial Transformer layer with Gated Linear Unit (GLU) activation function as described in:
  https://arxiv.org/pdf/1506.02025.pdf
  
  
  Parameters:
      in_channels (:obj:`int`):
          Input number of channels
      n_heads (:obj:`int`):
          Number of heads
      d_head (:obj:`int`):
          Hidden states dimension inside each head
      depth (:obj:`int`, *optional*, defaults to 1):
          Number of transformers block
      dropout (:obj:`float`, *optional*, defaults to 0.0):
          Dropout rate
      use_linear_projection (`bool`, defaults to `False`): tbd
      only_cross_attention (`bool`, defaults to `False`): tbd
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
      use_memory_efficient_attention (`bool`, *optional*, defaults to `False`):
          enable memory efficient attention https://arxiv.org/abs/2112.05682
      split_head_dim (`bool`, *optional*, defaults to `False`):
          Whether to split the head dimension into a new axis for the self-attention computation. In most cases,
          enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL.
code: def __call__(self, hidden_states, context, deterministic=True):
      batch, height, width, channels = hidden_states.shape
      residual = hidden_states
      hidden_states = self.norm(hidden_states)
      if self.use_linear_projection:
          hidden_states = hidden_states.reshape(
              batch, height * width, channels
          )
          hidden_states = self.proj_in(hidden_states)
      else:
          hidden_states = self.proj_in(hidden_states)
          hidden_states = hidden_states.reshape(
              batch, height * width, channels
          )
  
      for transformer_block in self.transformer_blocks:
          hidden_states = transformer_block(
              hidden_states, context, deterministic=deterministic
          )
  
      if self.use_linear_projection:
          hidden_states = self.proj_out(hidden_states)
          hidden_states = hidden_states.reshape(
              batch, height, width, channels
          )
      else:
          hidden_states = hidden_states.reshape(
              batch, height, width, channels
          )
          hidden_states = self.proj_out(hidden_states)
  
      hidden_states = hidden_states + residual
      return self.dropout_layer(
          hidden_states, deterministic=deterministic
      )
  
imports: import functools
  import math
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the FlaxTransformer2DModel obj's underlying member method code with example usage: obj.__call__(hidden_states, context, deterministic) and returns: self