class_docstr: Flax Convolutional based multi-head attention block for diffusion-based VAE.
  
  Parameters:
      channels (:obj:`int`):
          Input channels
      num_head_channels (:obj:`int`, *optional*, defaults to `None`):
          Number of attention heads
      num_groups (:obj:`int`, *optional*, defaults to `32`):
          The number of groups to use for group norm
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
class_docstr_tok: ["Flax","Convolutional","based","multi-head","attention","block","for","diffusion-based","VAE",".","Parameters",":","channels","(",":","obj",":","`","int","`",")",":","Input","channels","num_head_channels","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Number","of","attention","heads","num_groups","(",":","obj",":","`","int","`",",","*","optional","*",",","defaults","to","`","32","`",")",":","The","number","of","groups","to","use","for","group","norm","dtype","(",":","obj",":","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","jnp.float32",")",":","Parameters","`","dtype","`"]
code: def setup(self):
      self.num_heads = (
          self.channels // self.num_head_channels
          if self.num_head_channels is not None
          else 1
      )
  
      dense = partial(nn.Dense, self.channels, dtype=self.dtype)
  
      self.group_norm = nn.GroupNorm(
          num_groups=self.num_groups, epsilon=1e-6
      )
      self.query, self.key, self.value = dense(), dense(), dense()
      self.proj_attn = dense()
  
imports: import math
  from functools import partial
  from typing import Tuple
  import flax
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict
  from ..configuration_utils import ConfigMixin, flax_register_to_config
  from ..utils import BaseOutput
  from .modeling_flax_utils import FlaxModelMixin
label_desc: Use image node_id 1 for calling the FlaxAttentionBlock obj's underlying member method code with example usage: obj.setup() without return types