code: def add_explain_tokens(X, max_length, pool_idx):
      flat_x = tf.reshape(X[:, :, :1], [-1, 1])
      flat_pos = tf.minimum(
          X[:, :, 1:] + 1, max_length - 1
      )  # + 1 to offset for start token
      clf_tok = tf.gather(
          flat_x,
          tf.range(shape_list(X)[0], dtype=tf.int32) * max_length
          + pool_idx,
      )
      clf_tok_x_seq = tf.tile(
          tf.expand_dims(clf_tok, 1), [1, max_length, 1]
      )
      clf_tok_x_seq_w_pos = tf.concat((clf_tok_x_seq, flat_pos), -1)
      return tf.concat((X, clf_tok_x_seq_w_pos), 1)
  
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 16 for calling a global function with example usage: add_explain_tokens(X, max_length, pool_idx) and returns: tf