code: def _index_select_nd(
      source: torch.Tensor, index: torch.Tensor
  ) -> torch.Tensor:
      """
      Selects the message features from source corresponding to the atom or bond indices in index.
  
      Parameters
      ----------
      source: torch.Tensor
          A tensor of shape (num_bonds, hidden_size) containing message features.
      index: torch.Tensor
          A tensor of shape (num_atoms/num_bonds, max_num_bonds) containing the atom or bond indices to select from source.
  
      Returns
      ----------
      message_features: torch.Tensor
          A tensor of shape (num_atoms/num_bonds, max_num_bonds, hidden_size) containing the message features corresponding to the atoms/bonds specified in index.
      """
      index_size = index.size()  # (num_atoms/num_bonds, max_num_bonds)
      suffix_dim = source.size()[1:]  # (hidden_size,)
      final_size = (
          index_size + suffix_dim
      )  # (num_atoms/num_bonds, max_num_bonds, hidden_size)
  
      target = source.index_select(
          dim=0, index=index.view(-1)
      )  # (num_atoms/num_bonds * max_num_bonds, hidden_size)
      target = target.view(
          final_size
      )  # (num_atoms/num_bonds, max_num_bonds, hidden_size)
  
      return target
  
imports: from typing import List, Dict
  import numpy
  from scipy import stats
  from deepchem.models.torch_models.readout import GroverReadout
  from deepchem.models.torch_models.layers import SublayerConnection, PositionwiseFeedForward
label_desc: Use image node_id 1 for calling a global function with example usage: _index_select_nd(source, index) and returns: target