code: def put_atoms_in_cells(coords, cells, N, n_cells, ndim, k=5):
      """Place each atom into cells. O(N) runtime.
  
      Let N be the number of atoms.
  
      Parameters
      ----------
      coords: tf.Tensor
        (N, 3) shape.
      cells: tf.Tensor
        (n_cells, ndim) shape.
      N: int
        Number atoms
      ndim: int
        Dimensionality of input space
      k: int
        Number of nearest neighbors.
  
      Returns
      -------
      closest_atoms: tf.Tensor
        Of shape (n_cells, k, ndim)
      """
      n_cells = int(n_cells)
      # Tile both cells and coords to form arrays of size (n_cells*N, ndim)
      tiled_cells = tf.reshape(
          tf.tile(cells, (1, N)), (n_cells * N, ndim)
      )
      # TODO(rbharath): Change this for tf 1.0
      # n_cells tensors of shape (N, 1)
      tiled_cells = tf.split(tiled_cells, n_cells)
  
      # Shape (N*n_cells, 1) after tile
      tiled_coords = tf.tile(coords, (n_cells, 1))
      # List of n_cells tensors of shape (N, 1)
      tiled_coords = tf.split(tiled_coords, n_cells)
  
      # Lists of length n_cells
      coords_rel = [
          tf.cast(coords, tf.float32) - tf.cast(cells, tf.float32)
          for (coords, cells) in zip(tiled_coords, tiled_cells)
      ]
      coords_norm = [
          tf.reduce_sum(rel**2, axis=1) for rel in coords_rel
      ]
  
      # Lists of length n_cells
      # Get indices of k atoms closest to each cell point
      closest_inds = [tf.nn.top_k(norm, k=k)[1] for norm in coords_norm]
      # n_cells tensors of shape (k, ndim)
      closest_atoms = tf.stack(
          [tf.gather(coords, inds) for inds in closest_inds]
      )
      # Tensor of shape (n_cells, k)
      closest_inds = tf.stack(closest_inds)
  
      return closest_inds, closest_atoms
  
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 5 for calling a global function with example usage: put_atoms_in_cells(coords, cells, N, n_cells, ndim, k) and returns: closest_inds,  closest_atoms