code: def _batch_and_set(graphs, atom_x, bond_x, flatten):
      if flatten:
          graphs = [g for f in graphs for g in f]
      graph_batch = dgl.batch(graphs)
      graph_batch.ndata["x"] = atom_x
      graph_batch.edata.update(
          {
              "x": bond_x,
              "src_x": atom_x.new(
                  bond_x.shape[0], atom_x.shape[1]
              ).zero_(),
          }
      )
      return graph_batch
  
file_path: dgl/examples/pytorch/jtnn/jtnn/datautils.py
imports: import dgl
  import torch
  from dgl.data.utils import _get_dgl_url, download, extract_archive, get_download_dir
  from torch.utils.data import Dataset
  from .jtmpn import ATOM_FDIM, BOND_FDIM, mol2dgl_single
  from .mol_tree import Vocab
  from .mol_tree_nx import DGLMolTree
  from .mpn import mol2dgl_single
label_desc: Use image node_id 2 for calling the JTNNCollator obj's underlying member method code with example usage: obj._batch_and_set(graphs, atom_x, bond_x, flatten) and returns: graph_batch