code: def multi_head_dot_product_attention(
      scope: Scope,
      inputs_q: Array,
      inputs_kv: Array,
      bias: Optional[Array] = None,
      qkv_features: Optional[int] = None,
      out_features: Optional[int] = None,
      attn_fn: Callable = softmax_attn,
      batch_axes: Sequence[int] = (0,),
      num_heads: int = 1,
      dtype=jnp.float32,
      broadcast_dropout=False,
  ):
      if qkv_features is None:
          qkv_features = inputs_q.shape[-1]
      if out_features is None:
          out_features = inputs_q.shape[-1]
  
      attn_fn = partial(
          dot_product_attention,
          attn_fn=attn_fn,
          qkv_features=qkv_features // num_heads,
          out_features=out_features,
          dtype=dtype,
      )
      attn_fn = lift.vmap(
          attn_fn,
          in_axes=(None, None, None),
          out_axes=-2,
          axis_size=num_heads,
          variable_axes={"params": 0},
          split_rngs={"params": True, "dropout": not broadcast_dropout},
      )
      for axis in reversed(sorted(batch_axes)):
          attn_fn = lift.vmap(
              attn_fn,
              in_axes=(axis, axis, axis),
              out_axes=axis,
              variable_axes={"params": None},
              split_rngs={
                  "params": False,
                  "dropout": not broadcast_dropout,
              },
          )
  
      y = attn_fn(scope, inputs_q, inputs_kv, bias)
      return y.mean(axis=-2)
  
file_path: flax/tests/core/design/core_attention_test.py
imports: from functools import partial
  from typing import Callable, Optional, Sequence
  import jax
  from absl.testing import absltest
  from jax import lax, random
  from jax import numpy
  from flax.core import Array, Scope, init, lift, nn, unfreeze
label_desc: Use image node_id 5 for calling a global function with example usage: multi_head_dot_product_attention(scope, inputs_q, inputs_kv, bias, qkv_features, out_features, attn_fn, batch_axes, num_heads, dtype, broadcast_dropout) and returns: y