code: def compute_neighbor_list(
      coords, nbr_cutoff, N, M, n_cells, ndim=3, k=5
  ):
      """Computes a neighbor list from atom coordinates.
  
      Parameters
      ----------
      coords: tf.Tensor
        Shape (N, ndim)
      N: int
        Max number atoms
      M: int
        Max number neighbors
      ndim: int
        Dimensionality of space.
      k: int
        Number of nearest neighbors to pull down.
  
      Returns
      -------
      nbr_list: tf.Tensor
        Shape (N, M) of atom indices
      """
      start = tf.cast(tf.reduce_min(coords), tf.int32)
      stop = tf.cast(tf.reduce_max(coords), tf.int32)
      cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
      # Associate each atom with cell it belongs to. O(N*n_cells)
      # Shape (n_cells, k)
      atoms_in_cells, _ = put_atoms_in_cells(
          coords, cells, N, n_cells, ndim, k
      )
      # Shape (N, 1)
      cells_for_atoms = get_cells_for_atoms(
          coords, cells, N, n_cells, ndim
      )
  
      # Associate each cell with its neighbor cells. Assumes periodic boundary
      # conditions, so does wrapround. O(constant)
      # Shape (n_cells, 26)
      neighbor_cells = compute_neighbor_cells(cells, ndim, n_cells)
  
      # Shape (N, 26)
      neighbor_cells = tf.squeeze(
          tf.gather(neighbor_cells, cells_for_atoms)
      )
  
      # coords of shape (N, ndim)
      # Shape (N, 26, k, ndim)
      tiled_coords = tf.tile(
          tf.reshape(coords, (N, 1, 1, ndim)), (1, 26, k, 1)
      )
  
      # Shape (N, 26, k)
      nbr_inds = tf.gather(atoms_in_cells, neighbor_cells)
  
      # Shape (N, 26, k)
      atoms_in_nbr_cells = tf.gather(atoms_in_cells, neighbor_cells)
  
      # Shape (N, 26, k, ndim)
      nbr_coords = tf.gather(coords, atoms_in_nbr_cells)
  
      # For smaller systems especially, the periodic boundary conditions can
      # result in neighboring cells being seen multiple times. Maybe use tf.unique to
      # make sure duplicate neighbors are ignored?
  
      # TODO(rbharath): How does distance need to be modified here to
      # account for periodic boundary conditions?
      # Shape (N, 26, k)
      dists = tf.reduce_sum((tiled_coords - nbr_coords) ** 2, axis=3)
  
      # Shape (N, 26*k)
      dists = tf.reshape(dists, [N, -1])
  
      # TODO(rbharath): This will cause an issue with duplicates!
      # Shape (N, M)
      closest_nbr_locs = tf.nn.top_k(dists, k=M)[1]
  
      # N elts of size (M,) each
      split_closest_nbr_locs = [
          tf.squeeze(locs) for locs in tf.split(closest_nbr_locs, N)
      ]
  
      # Shape (N, 26*k)
      nbr_inds = tf.reshape(nbr_inds, [N, -1])
  
      # N elts of size (26*k,) each
      split_nbr_inds = [
          tf.squeeze(split) for split in tf.split(nbr_inds, N)
      ]
  
      # N elts of size (M,) each
      neighbor_list = [
          tf.gather(nbr_inds, closest_nbr_locs)
          for (nbr_inds, closest_nbr_locs) in zip(
              split_nbr_inds, split_closest_nbr_locs
          )
      ]
  
      # Shape (N, M)
      neighbor_list = tf.stack(neighbor_list)
  
      return neighbor_list
  
imports: from __future__ import print_function
  from __future__ import division
  from __future__ import unicode_literals
  import warnings
  import numpy
  import tensorflow
  from deepchem.models import Model
  from deepchem.nn import model_ops
  import deepchem.utils.rdkit_util
label_desc: Use image node_id 1 for calling a global function with example usage: compute_neighbor_list(coords, nbr_cutoff, N, M, n_cells, ndim, k) and returns: neighbor_list