code: def compute_closest_neighbors(
      coords,
      cells,
      atoms_in_cells,
      neighbor_cells,
      N,
      n_cells,
      ndim=3,
      k=5,
  ):
      """Computes nearest neighbors from neighboring cells.
  
      TODO(rbharath): Make this pass test
  
      Parameters
      ---------
      atoms_in_cells: list
        Of length n_cells. Each entry tensor of shape (k, ndim)
      neighbor_cells: tf.Tensor
        Of shape (n_cells, 26).
      N: int
        Number atoms
      """
      n_cells = int(n_cells)
      # Tensor of shape (n_cells, k, ndim)
      # atoms_in_cells = tf.stack(atoms_in_cells)
  
      cells_for_atoms = get_cells_for_atoms(
          coords, cells, N, n_cells, ndim
      )
      all_closest = []
      for atom in range(N):
          atom_vec = coords[atom]
          cell = cells_for_atoms[atom]
          nbr_inds = tf.gather(neighbor_cells, tf.cast(cell, tf.int32))
          # Tensor of shape (26, k, ndim)
          nbr_atoms = tf.gather(atoms_in_cells, nbr_inds)
          # Reshape to (26*k, ndim)
          nbr_atoms = tf.reshape(nbr_atoms, (-1, 3))
          # Subtract out atom vector. Still of shape (26*k, ndim) due to broadcast.
          nbr_atoms = nbr_atoms - atom_vec
          # Dists of shape (26*k, 1)
          nbr_dists = tf.reduce_sum(nbr_atoms**2, axis=1)
          # Of shape (k, ndim)
          closest_inds = tf.nn.top_k(nbr_dists, k=k)[1]
          all_closest.append(closest_inds)
      return all_closest
  
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 3 for calling a global function with example usage: compute_closest_neighbors(coords, cells, atoms_in_cells, neighbor_cells, N, n_cells, ndim, k) and returns: all_closest