code: def attn_weights(
      q, k, v, scale=False, mask=True, explain=False, lengths=None
  ):
      w = tf.matmul(q, k)
  
      if scale:
          n_state = shape_list(v)[-1]
          w = w * tf.math.rsqrt(tf.cast(n_state, tf.float32))
  
      if mask:
          if explain:
              w = explain_mask_attn_weights(w)
          else:
              w = mask_attn_weights(w)
      elif lengths is not None:
          # at least mask pad tokens
          w = mask_pad(w, lengths=lengths)
  
      w = tf.nn.softmax(w)
      return w
  
imports: import functools
  import numpy
  import tensorflow
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.nn.activations import act_fns
  from finetune.nn.nn_utils import dropout, norm
label_desc: Use image node_id 4 for calling a global function with example usage: attn_weights(q, k, v, scale, mask, explain, lengths) and returns: w