code: def attn(x, scope, n_state, *, past, hparams):
      assert x.shape.ndims == 3  # Should be [batch, sequence, features]
      assert n_state % hparams.n_head == 0
      if past is not None:
          assert (
              past.shape.ndims == 5
          )  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]
  
      def split_heads(x):
          # From [batch, sequence, features] to [batch, heads, sequence, features]
          return tf.transpose(
              a=split_states(x, hparams.n_head), perm=[0, 2, 1, 3]
          )
  
      def merge_heads(x):
          # Reverse of split_heads
          return merge_states(tf.transpose(a=x, perm=[0, 2, 1, 3]))
  
      def mask_attn_weights(w):
          # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
          _, _, nd, ns = shape_list(w)
          b = attention_mask(nd, ns, dtype=w.dtype)
          b = tf.reshape(b, [1, 1, nd, ns])
          w = w * b - tf.cast(1e10, w.dtype) * (1 - b)
          return w
  
      def multihead_attn(q, k, v):
          # q, k, v have shape [batch, heads, sequence, features]
          w = tf.matmul(q, k, transpose_b=True)
          w = w * tf.math.rsqrt(
              tf.cast(get_value(v.shape[-1]), w.dtype)
          )
  
          w = mask_attn_weights(w)
          w = softmax(w)
          a = tf.matmul(w, v)
          return a
  
      with tf.compat.v1.variable_scope(scope):
          c = conv1d(x, "c_attn", n_state * 3)
          q, k, v = map(split_heads, tf.split(c, 3, axis=2))
          present = tf.stack([k, v], axis=1)
          if past is not None:
              pk, pv = tf.unstack(past, axis=1)
              k = tf.concat([pk, k], axis=-2)
              v = tf.concat([pv, v], axis=-2)
          a = multihead_attn(q, k, v)
          a = merge_heads(a)
          a = conv1d(a, "c_proj", n_state)
          return a, present
  
imports: import numpy
  import tensorflow
label_desc: Use image node_id 10 for calling a global function with example usage: attn(x, scope, n_state) and returns: tf, merge_states, w, a, a,  present