def _nested_concatenate(input_graphs, field_name, axis):
  """Concatenates a possibly nested feature field of a list of input graphs."""
  features_list = [getattr(gr, field_name) for gr in input_graphs
                   if getattr(gr, field_name) is not None]
  if not features_list:
    return None

  if len(features_list) < len(input_graphs):
    raise ValueError(
        "All graphs or no graphs must contain {} features.".format(field_name))

  name = "concat_" + field_name
  return tree.map_structure(lambda *x: tf.concat(x, axis, name), *features_list)
Concatenates a possibly nested feature field of a list of input graphs.
graph_nets/utils_tf.py