code: def explain_mask_attn_weights(w):
      # w is [batch, heads, n, n]
      # lengths is [batch]
      batch, _, _, n = shape_list(w)
      seq = n // 2
      main_mask = tf.linalg.band_part(tf.ones([seq, seq]), -1, 0)
      top = tf.expand_dims(
          tf.concat((main_mask, tf.zeros([seq, seq])), 1), 0
      )  # 1, seq, 2 * seq
      clf_to_clf_mask = tf.eye(seq)
      bottom = tf.expand_dims(
          tf.concat((main_mask, clf_to_clf_mask), 1), 0
      )  # 1, seq, 2 * seq
      m = tf.concat((top, bottom), 1)
      b = tf.reshape(m, [1, 1, n, n])
      w = w * b + -1e9 * (1 - b)
      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 3 for calling a global function with example usage: explain_mask_attn_weights(w) and returns: w