code: def get_cells_for_atoms(coords, cells, N, n_cells, ndim=3):
      """Compute the cells each atom belongs to.
  
      Parameters
      ----------
      coords: tf.Tensor
        Shape (N, ndim)
      cells: tf.Tensor
        (box_size**ndim, ndim) shape.
      Returns
      -------
      cells_for_atoms: tf.Tensor
        Shape (N, 1)
      """
      n_cells = int(n_cells)
      # Tile both cells and coords to form arrays of size (n_cells*N, ndim)
      tiled_cells = tf.tile(cells, (N, 1))
      # N tensors of shape (n_cells, 1)
      tiled_cells = tf.split(tiled_cells, N)
  
      # Shape (N*n_cells, 1) after tile
      tiled_coords = tf.reshape(
          tf.tile(coords, (1, n_cells)), (n_cells * N, ndim)
      )
      # List of N tensors of shape (n_cells, 1)
      tiled_coords = tf.split(tiled_coords, N)
  
      # Lists of length N
      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=1)[1] for norm in coords_norm
      ]
  
      # TODO(rbharath): tf.stack for tf 1.0
      return tf.stack(closest_inds)
  
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 2 for calling a global function with example usage: get_cells_for_atoms(coords, cells, N, n_cells, ndim) and returns: tf