code: def gen_mask(
      batch_ids, mask_type="bidi", query_len=None, pad_value=0
  ):
      if query_len is None:
          query_len = batch_ids.shape[1]
      if mask_type != "empty":
          mask = (batch_ids != pad_value).astype(np.float32)
          mask = np.tile(np.expand_dims(mask, 1), [1, query_len, 1])
          if mask_type == "causal":
              assert query_len == batch_ids.shape[1]
              mask = np.tril(mask)
          elif mask_type == "causal_without_diag":
              assert query_len == batch_ids.shape[1]
              mask = np.tril(mask, -1)
          elif mask_type == "diag":
              assert query_len == batch_ids.shape[1]
              # import pdb; pdb.set_trace()
              mask = np.stack([np.diag(np.diag(m)) for m in mask], 0)
  
      else:
          mask_type == "empty"
          mask = np.zeros_like(batch_ids).astype(np.float32)
          mask = np.tile(np.expand_dims(mask, 1), [1, query_len, 1])
      return mask
  
imports: from copy import deepcopy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: gen_mask(batch_ids, mask_type, query_len, pad_value) and returns: mask