code: def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):
      qkv_features = self.qkv_features or inputs_q.shape[-1]
      out_features = self.out_features or inputs_q.shape[-1]
  
      # Now, vmap attn.__call__ along heads and spatial dims.
      Attn = concise_vmap(
          DotProductAttention,
          (None, None, None),
          -2,
          param=(0, True),
          dropout=(None, not self.broadcast_dropout),
          axis_size=self.num_heads,
      )
      for axis in reversed(sorted(self.batch_axes)):
          Attn = concise_vmap(
              Attn,
              (axis, axis, axis),
              axis,
              param=(None, False),
              dropout=(None, not self.broadcast_dropout),
          )
  
      attn = Attn(
          attn_module=self.attn_module,
          qkv_features=qkv_features // self.num_heads,
          out_features=out_features,
      )
  
      # evaluate multi-headed-attention.
      y = attn(inputs_q, inputs_kv, bias)
      return y.mean(axis=-2)
  
imports: import functools
  from pprint import pprint
  from typing import Any, Callable, Optional, Sequence
  from flax.core.frozen_dict import unfreeze
  from flax.linen import initializers
  from flax.linen import Module, compact, vmap
  from flax.linen.linear import PrecisionLike
  import jax
  from jax import lax, numpy, random
label_desc: Use image node_id 1 for calling the MultiHeadDotProductAttention obj's underlying member method code with example usage: obj.__call__(inputs_q, inputs_kv, bias, dtype) and returns: y