code: def DotProductAttention(
      queries,
      keys,
      values,
      pos_emb,
      context_bias,
      location_bias,
      mask,
      separate_cls,
      dropout,
      mode,
      rng,
  ):
      """Computes new activations via masked attention-weighted sum of values.
  
      Args:
        queries: Per-head activations representing attention queries.
        keys: Per-head activations representing attention keys.
        values: Per-head activations to be combined by computed attention weights.
        pos_emb: Per-head activations representing positional embeddings.
        context_bias: Global context bias from Transformer XL's attention.
        location_bias: Global location bias from Transformer XL's attention.
        mask: Mask that distinguishes positions with real content vs. padding.
        separate_cls: True/False if we separate_cls in calculations.
        dropout: Probabilistic rate for dropout applied to attention strengths
          (based on query-key pairs) before applying them to values.
        mode: One of `'train'`, `'eval'`, or `'predict'`.
        rng: Single-use random number generator (JAX PRNG key).
  
      Returns:
        Per-head activations resulting from masked per-head attention-weighted
        sum of per-head values.
  
      This function is the core of the attention mechanism. It:
        - computes per-head attention weights from per-head `queries` and `keys`,
        - applies `mask` to screen out positions that come from padding tokens,
        - optionally applies dropout to attention weights, and
        - uses attention weights to combine per-head `values` vectors.
      """
      d_feature = queries.shape[-1]
      keys_len, queries_len = keys.shape[-2], queries.shape[-2]
      funnel_factor, is_upsampling = calc_funnel_ratio(
          keys_len, queries_len
      )
  
      ac = jnp.einsum("bnid,bnjd->bnij", queries + context_bias, keys)
      bd = jnp.einsum(
          "bnid,jnd->bnij", queries + location_bias, pos_emb
      )
      bd = _fast_matrix_shift(bd, funnel_factor, is_upsampling)
  
      if separate_cls:
          # Masking out location part of attention for cls token
          bd = bd.at[:, :, :, 0].set(0)
          bd = bd.at[:, :, 0, :].set(0)
  
      dots = (ac + bd) / jnp.sqrt(d_feature)
      if mask is not None:
          dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
      # Softmax.
      dots = jnp.exp(
          dots - fastmath.logsumexp(dots, axis=-1, keepdims=True)
      )
      if dropout >= 1.0:
          raise ValueError("Dropout rates must be lower than 1.")
      if dropout is not None and dropout > 0.0 and mode == "train":
          keep = fastmath.random.bernoulli(
              rng, 1.0 - dropout, dots.shape
          )
          dots = jnp.where(
              keep, dots / (1.0 - dropout), jnp.zeros_like(dots)
          )
      out = jnp.matmul(dots, values)
      out = out.astype(jnp.float32)
      dots = dots.astype(jnp.float32)
      return out, dots
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 5 for calling a global function with example usage: DotProductAttention(queries, keys, values, pos_emb, context_bias, location_bias, mask, separate_cls, dropout, mode, rng) and returns: out,  dots