code: def collect_probs(probs, positions):
      """
      Collect Probabilities.
  
      Reference:
      https://github.com/zhiguowang/BiMPM/blob/master/src/layer_utils.py#L128-L140
      :param probs: [batch_size, chunks_size]
      :param positions: [batch_size, pair_size]
      :return: [batch_size, pair_size]
      """
      batch_size = tf.shape(probs)[0]
      pair_size = tf.shape(positions)[1]
      # shape (batch_size)
      batch_nums = K.arange(0, batch_size)
      # [batch_size, 1]
      batch_nums = tf.reshape(batch_nums, shape=[-1, 1])
      # [batch_size, pair_size]
      batch_nums = K.tile(batch_nums, [1, pair_size])
  
      # shape (batch_size, pair_size, 2)
      # Alert: to solve error message
      positions = tf.cast(positions, tf.int32)
      indices = tf.stack([batch_nums, positions], axis=2)
  
      pair_probs = tf.gather_nd(probs, indices)
      # pair_probs = tf.reshape(pair_probs, shape=[batch_size, pair_size])
      return pair_probs
  
imports: import tensorflow
  from keras import backend
  from keras.engine import Layer
  from matchzoo.contrib.layers.attention_layer import AttentionLayer
label_desc: Use image node_id 4 for calling a global function with example usage: collect_probs(probs, positions) and returns: pair_probs