code: def batched_gather(data, inds, dim=0, no_batch_dims=0):
      ranges = []
      for i, s in enumerate(data.shape[:no_batch_dims]):
          r = torch.arange(s)
          r = r.view(
              *(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1)))
          )
          ranges.append(r)
  
      remaining_dims = [
          slice(None) for _ in range(len(data.shape) - no_batch_dims)
      ]
      remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds
      ranges.extend(remaining_dims)
      # Matt note: Editing this to get around the behaviour of using a list as an array index changing
      # in recent Numpy versions
      return data[tuple(ranges)]
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/esm/openfold_utils/tensor_utils.py
imports: from functools import partial
  from typing import List
  import torch
  import torch.nn
label_desc: Use image node_id 8 for calling a global function with example usage: batched_gather(data, inds, dim, no_batch_dims) and returns: data