code: def get_seg_mask_from_token_idx(inputs, seg_idx):
      """
      get padding mask from the input token idx
      inputs: [batch_size, time_steps]
      mask: [batch_size, time_steps]
      """
      seg_mask = tf.cast(tf.math.equal(inputs, seg_idx), tf.int32)
      return seg_mask
  
file_path: delta/delta/layers/utils.py
imports: import math
  import delta.compat
  from tensorflow.python.keras import backend
  import delta.utils
label_desc: Use image node_id 8 for calling a global function with example usage: get_seg_mask_from_token_idx(inputs, seg_idx) and returns: seg_mask