code: def _per_head_attention(
      queries, keys, values, mask, dropout, mode, rng
  ):
      """Computes new per-head activations via scaled dot-product attention.
  
      This function is the core of the attention mechanism. Given per-head
      ``queries`` (Q), ``keys`` (K), ``values`` (V), and ``mask``, it:
  
        - computes the scaled dot product of each Q-K pair;
        - applies ``mask`` to screen out positions that come from padding tokens
          (indicated by 0 value);
        - [in ``'train'`` mode] applies dropout to Q-K dot products;
        - computes Q-K attention strengths using a per-query softmax of the Q-K dot
          products; and
        - for each query position, combines V vectors according to the Q-K
          attention strengths.
  
      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 strengths.
        mask: Mask that distinguishes positions with real content vs. padding.
        dropout: Probababilistic rate for attention dropout, which overrides
            (sets to zero) some attention strengths derived from query-key
            matching. As a result, on a given forward pass, some value vectors
            don't contribute to the output, analogous to how regular dropout can
            cause some node activations to be ignored. Applies only in ``'train'``
            mode.
        mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
        rng: Single-use random number generator (JAX PRNG key).
  
      Returns:
        Tuple of (activations, attn_strengths), where activations are new per-head
        activation vectors and attn_strengths is a matrix of per-head attention
        strengths.
      """
      if dropout >= 1.0:
          raise ValueError(
              f"Dropout rate ({dropout}) must be lower than 1."
          )
  
      d_feature = queries.shape[-1]
  
      dots = jnp.matmul(queries, jnp.swapaxes(keys, -1, -2)) / jnp.sqrt(
          d_feature
      )
      if mask is not None:
          dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
      attn_strengths = jnp.exp(
          dots - fastmath.logsumexp(dots, axis=-1, keepdims=True)
      )
      if dropout is not None and dropout > 0.0 and mode == "train":
          keep = fastmath.random.bernoulli(
              rng, 1.0 - dropout, attn_strengths.shape
          )
          attn_strengths = jnp.where(
              keep,
              attn_strengths / (1.0 - dropout),
              jnp.zeros_like(attn_strengths),
          )
      activations = jnp.matmul(attn_strengths, values).astype(
          jnp.float32
      )
      attn_strengths = attn_strengths.astype(jnp.float32)
      return activations, attn_strengths
  
imports: import numpy
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
  from trax.layers.research import sparsity
label_desc: Use image node_id 3 for calling a global function with example usage: _per_head_attention(queries, keys, values, mask, dropout, mode, rng) and returns: activations,  attn_strengths