code: def batched_index_select(values, indices, dim=1):
      value_dims = values.shape[(dim + 1) :]
      values_shape, indices_shape = map(
          lambda t: list(t.shape), (values, indices)
      )
      indices = indices[(..., *((None,) * len(value_dims)))]
      indices = indices.expand(
          *((-1,) * len(indices_shape)), *value_dims
      )
      value_expand_len = len(indices_shape) - (dim + 1)
      values = values[
          (*((slice(None),) * dim), *((None,) * value_expand_len), ...)
      ]
  
      value_expand_shape = [-1] * len(values.shape)
      expand_slice = slice(dim, (dim + value_expand_len))
      value_expand_shape[expand_slice] = indices.shape[expand_slice]
      values = values.expand(*value_expand_shape)
  
      dim += value_expand_len
      return values.gather(dim, indices)
  
imports: import torch
  import torch.nn.functional
  from torch.nn.utils.rnn import pad_sequence
  from torch import nn, einsum
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 5 for calling a global function with example usage: batched_index_select(values, indices, dim) and returns: values