code: def attention_mask(nd, ns, *, dtype):
      """1's in the lower triangle, counting from the lower right corner.
      Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
      """
      i = tf.range(nd)[:, None]
      j = tf.range(ns)
      m = i >= j - ns + nd
      return tf.cast(m, dtype)
  
file_path: finetune/finetune/base_models/gpt2/featurizer.py
imports: import functools
  import tensorflow
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.nn.activations import gelu
  from finetune.base_models.gpt.featurizer import norm, dropout, get_pos_values
label_desc: Use image node_id 5 for calling a global function with example usage: attention_mask(nd, ns) and returns: tf