code: def _batched_index_select(
      target: torch.Tensor,
      indices: torch.LongTensor,
      dim: int = 1,
      flattened_indices: Optional[torch.LongTensor] = None,
  ) -> torch.Tensor:
      # Input:
      #   target: (batch_size, seq_len, dim)
      #   indices: (batch_size, num_indices)
      # Returns:
      #   (batch_size, num_indices, dim)
      unidim = False
      if len(target.size()) == len(indices.size()):
          # (batch_size, sequence_length) -> (batch_size, sequence_length, 1)
          unidim = True
          target = target.unsqueeze(-1)
  
      target_size = target.size()
      indices_size = indices.size()
  
      target = target.reshape(
          math.prod([*target_size[:dim]]), *target_size[dim:]
      )
      indices = indices.view(
          math.prod([*indices_size[:dim]]), *indices_size[dim:]
      )
  
      if flattened_indices is None:
          # Shape: (batch_size * d_1 * ... * d_n)
          flattened_indices = flatten_and_batch_shift_indices(
              indices, target.size(1)
          )
  
      # Shape: (batch_size * sequence_length, embedding_size)
      flattened_target = target.reshape(-1, *target_size[dim + 1 :])
      # Shape: (batch_size * d_1 * ... * d_n, embedding_size)
      flattened_selected = flattened_target.index_select(
          0, flattened_indices
      )
      selected_shape = (
          list(indices_size)
          if unidim
          else (list(indices_size) + list(target_size[dim + 1 :]))
      )
      # Shape: (batch_size, d_1, ..., d_n, embedding_size)
      selected_targets = flattened_selected.reshape(*selected_shape)
      return selected_targets
  
imports: import sys
  import os
  from os import makedirs
  from os.path import join
  import numpy
  import pyhocon
  import math
  import logging
  import random
  from typing import Optional, Tuple, Any, Dict, Iterable, Union
  import torch
  import torch.nn.functional
  from torch import nn
  from transformers.activations import ACT2FN
label_desc: Use image node_id 13 for calling a global function with example usage: _batched_index_select(target, indices, dim, flattened_indices) and returns: selected_targets