def gather(src, indices):
  """A partial implementation of PyTorch gather, for dim=1.

  Given a src tensor of variable size, this function gathers multiple elements
  along the 2nd dimension (dim=1), for every batch in the 1st dimension (dim=0).

  Args:
    src: <T> [batch_size, d_1, ..., d_n]
    indices: <int32> [batch_size, num_indices]

  Returns:
    output: <T> [batch_size, num_indices, d_2, ..., d_n]
  """
  # Hereby begins the painful dance of tf.gather_nd.
  batch_size = shape(indices, 0)
  num_indices = shape(indices, 1)

  # [batch_size, 1]
  batch_idx = tf.reshape(tf.range(batch_size), [batch_size, 1])

  # [batch_size, num_indices]
  batch_idx = tf.tile(batch_idx, [1, num_indices])

  # [batch_size * num_indices]
  batch_idx = tf.reshape(batch_idx, [-1])

  # [batch_size * num_indices]
  flat_indices = tf.reshape(indices, [-1])

  # [batch_size * num_indices, 2]
  idx = tf.stack([batch_idx, flat_indices], 1)

  # [batch_size * num_indices, d_2, ..., d_n]
  output = tf.gather_nd(src, idx)

  # [batch_size, num_indices, d_2, ..., d_n]
  output = tf.reshape(output, [batch_size, num_indices] +
                      [shape(output, i) for i in range(1, len(output.shape))])

  return output
A partial implementation of PyTorch gather, for dim=1.

  Given a src tensor of variable size, this function gathers multiple elements
  along the 2nd dimension (dim=1), for every batch in the 1st dimension (dim=0).

  Args:
    src: <T> [batch_size, d_1, ..., d_n]
    indices: <int32> [batch_size, num_indices]

  Returns:
    output: <T> [batch_size, num_indices, d_2, ..., d_n]
  
language/capwap/utils/tensor_utils.py