def concat(input_graphs, axis, name="graph_concat"):
  """Returns an op that concatenates graphs along a given axis.

  In all cases, the NODES, EDGES and GLOBALS dimension are concatenated
  along `axis` (if a fields is `None`, the concatenation is just a `None`).
  If `axis` == 0, then the graphs are concatenated along the (underlying) batch
  dimension, i.e. the RECEIVERS, SENDERS, N_NODE and N_EDGE fields of the tuples
  are also concatenated together.
  If `axis` != 0, then there is an underlying assumption that the receivers,
  SENDERS, N_NODE and N_EDGE fields of the graphs in `values` should all match,
  but this is not checked by this op.
  The graphs in `input_graphs` should have the same set of keys for which the
  corresponding fields is not `None`.

  Args:
    input_graphs: A list of `graphs.GraphsTuple` objects containing `Tensor`s
      and satisfying the constraints outlined above.
    axis: An axis to concatenate on.
    name: (string, optional) A name for the operation.

  Returns: An op that returns the concatenated graphs.

  Raises:
    ValueError: If `values` is an empty list, or if the fields which are `None`
      in `input_graphs` are not the same for all the graphs.
  """
  if not input_graphs:
    raise ValueError("List argument `input_graphs` is empty")
  utils_np._check_valid_sets_of_keys([gr._asdict() for gr in input_graphs])  # pylint: disable=protected-access
  if len(input_graphs) == 1:
    return input_graphs[0]

  with tf.name_scope(name):
    nodes = _nested_concatenate(input_graphs, NODES, axis)
    edges = _nested_concatenate(input_graphs, EDGES, axis)
    globals_ = _nested_concatenate(input_graphs, GLOBALS, axis)

    output = input_graphs[0].replace(nodes=nodes, edges=edges, globals=globals_)
    if axis != 0:
      return output
    n_node_per_tuple = tf.stack(
        [tf.reduce_sum(gr.n_node) for gr in input_graphs])
    n_edge_per_tuple = tf.stack(
        [tf.reduce_sum(gr.n_edge) for gr in input_graphs])
    offsets = _compute_stacked_offsets(n_node_per_tuple, n_edge_per_tuple)
    n_node = tf.concat(
        [gr.n_node for gr in input_graphs], axis=0, name="concat_n_node")
    n_edge = tf.concat(
        [gr.n_edge for gr in input_graphs], axis=0, name="concat_n_edge")
    receivers = [
        gr.receivers for gr in input_graphs if gr.receivers is not None
    ]
    receivers = receivers or None
    if receivers:
      receivers = tf.concat(receivers, axis, name="concat_receivers") + offsets
    senders = [gr.senders for gr in input_graphs if gr.senders is not None]
    senders = senders or None
    if senders:
      senders = tf.concat(senders, axis, name="concat_senders") + offsets
    return output.replace(
        receivers=receivers, senders=senders, n_node=n_node, n_edge=n_edge)
Returns an op that concatenates graphs along a given axis.

  In all cases, the NODES, EDGES and GLOBALS dimension are concatenated
  along `axis` (if a fields is `None`, the concatenation is just a `None`).
  If `axis` == 0, then the graphs are concatenated along the (underlying) batch
  dimension, i.e. the RECEIVERS, SENDERS, N_NODE and N_EDGE fields of the tuples
  are also concatenated together.
  If `axis` != 0, then there is an underlying assumption that the receivers,
  SENDERS, N_NODE and N_EDGE fields of the graphs in `values` should all match,
  but this is not checked by this op.
  The graphs in `input_graphs` should have the same set of keys for which the
  corresponding fields is not `None`.

  Args:
    input_graphs: A list of `graphs.GraphsTuple` objects containing `Tensor`s
      and satisfying the constraints outlined above.
    axis: An axis to concatenate on.
    name: (string, optional) A name for the operation.

  Returns: An op that returns the concatenated graphs.

  Raises:
    ValueError: If `values` is an empty list, or if the fields which are `None`
      in `input_graphs` are not the same for all the graphs.
  
graph_nets/utils_tf.py