code: def gather(data, index_2d):
      """
      Gather data alongs the second dim. Assume data is 3d with shape (batch_size, dim1, dim2),
      and index_2d's shape is (batch_size, dim1).
      output[i][j] = data[i][index_2d[i][j]]
  
      This function does not require data, output, or index_2d having the same shape, which
       is mandated by torch.gather.
      """
      batch_size = data.shape[0]
      data_dim = data.shape[2]
      index_len = index_2d.shape[1]
      device = data.device
      res = data[
          torch.arange(batch_size, device=device).repeat_interleave(
              # index_len has to be moved to the device explicitly, otherwise
              # error will throw during jit.trace
              torch.tensor([index_len], device=device)
          ),
          index_2d.flatten(),
      ].view(batch_size, index_len, data_dim)
      return res
  
file_path: ReAgent/reagent/core/torch_utils.py
imports: from io import BytesIO
  from typing import Dict, List
  import numpy
  import torch
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
label_desc: Use image node_id 7 for calling a global function with example usage: gather(data, index_2d) and returns: res