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 = np.arange(nd)[:, None]
      j = np.arange(ns)
      m = i >= j - ns + nd
      return np.array(m, dtype=dtype)
  
imports: import json
  import os
  import sys
  import jax.numpy
  import numpy
  import objax
  from objax.variable import TrainVar
label_desc: Use image node_id 5 for calling a global function with example usage: attention_mask(nd, ns) and returns: np