code: def gather_neighbors(X, nbr_indices, B, N, M, d):
      """Gathers the neighbor subsets of the atoms in X.
  
      B = batch_size, N = max_num_atoms, M = max_num_neighbors, d = num_features
  
      Parameters
      ----------
      X: tf.Tensor of shape (B, N, d)
        Coordinates/features tensor.
      atom_indices: tf.Tensor of shape (B, M)
        Neighbor list for single atom.
  
      Returns
      -------
      neighbors: tf.Tensor of shape (B, M, d)
        Neighbor coordinates/features tensor for single atom.
  
      """
  
      example_tensors = tf.unstack(X, axis=0)
      example_nbrs = tf.unstack(nbr_indices, axis=0)
      all_nbr_coords = []
      for example, (example_tensor, example_nbr) in enumerate(
          zip(example_tensors, example_nbrs)
      ):
          nbr_coords = tf.gather(example_tensor, example_nbr)
          all_nbr_coords.append(nbr_coords)
      neighbors = tf.stack(all_nbr_coords)
      return neighbors
  
file_path: deepchem/contrib/atomicconv/models/atomicnet_ops.py
imports: from __future__ import print_function
  from __future__ import division
  from __future__ import unicode_literals
  import numpy
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: gather_neighbors(X, nbr_indices, B, N, M, d) and returns: neighbors