def pad_graphs_tuple(graphs_tuple,
                     pad_nodes_to,
                     pad_edges_to,
                     pad_graphs_to,
                     experimental_unconnected_padding_edges=False):
  """Pads a `graphs.GraphsTuple` to fixed number of nodes, edges and graphs.

  The Graph Nets library treat zeros as valid parts of a graph.GraphsTuple, so
  special padding is required in order to preserve the computation. This
  method does so by adding a 'dummy' graph to the batch so that additional
  nodes/edges can't interfere with the valid graph.

  Args:
    graphs_tuple: `graphs.GraphsTuple` batch of graphs.
    pad_nodes_to: the size to pad node determined features to.
    pad_edges_to: the size to pad edge determined features to.
    pad_graphs_to: the size to pad graph determined features to.
    experimental_unconnected_padding_edges: Experimental feature to prevent nans
      in the padding graph. DISCLAIMER: This feature is extremly experimental,
      and setting it to `True` is not recommened unless strictly necessary, and
      understanding the implications.

      If `True`, the padding graph will have `senders` and `receivers` for
      the padding edges reference a node which does not exist (one beyond the
      size of `nodes`).

      This feature can be used to prevent any broadcasting/aggregation ops
      between edges and nodes for the padding graph. The reason is that the
      sum aggregations in the padding graph, which has a single node with a
      very large number of self-edges, sometimes lead to infs or nans,
      which may contaminate the gradients of the other valid graphs in the batch
      with nans (even if masked out of the loss: this is related to the
      `tf.where` issue.).

      This approach relies on several numerical hacks that do not work on CPU,
      but work on GPU and TPU (as covered by our tests):

      * `tf.gather` returns zeros when the index is beyond the boundaries. From
         https://www.tensorflow.org/api_docs/python/tf/gather
           "Note that on CPU, if an out of bound index is found, an error is
           returned. On GPU, if an out of bound index is found, a 0 is stored
           in the corresponding output value."
      * `tf.unsorted_segment_sum` drops values for negative indices. From
         https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum
           "If the given segment ID is negative, the value is dropped and
           will not be added to the sum of the segment."
        We have seen empirically that it also ignores values with indices equal
        or larger than `num_segments`. While this behavior is tested in our
        library, we cannot guarantee that it will work in the future for all
        unsorted_segment ops, so use at your own risk.

      This fixes the appearance of nans in the node-wise edge aggregation. The
      appearance of `nan`s is less likely in the global aggregation because in
      the worst case, the number of nodes/edges on the padding graph is not
      typically much larger than the number of nodes/edges in other graphs in
      the dataset.

      A less hacky approach (but more expensive, and requiring modifying model
      code) to prevent nan's appearing in the padding graph, is by masking out
      the graph features before they are aggregated, although for convenience
      we usually find that it is enough to do it after each message passing
      layer. E.g.:

      ```
      graphs_tuple_size = get_graphs_tuple_size(graphs_tuple)

      padded_graphs_tuple = pad_graphs_tuple(graphs_tuple, ...)

      graphs_mask = get_mask(graphs_tuple_size.num_graphs, pad_graphs_to)
      nodes_mask = get_mask(graphs_tuple_size.num_nodes, pad_nodes_to)
      edges_mask = get_mask(graphs_tuple_size.num_edges, pad_edges_to)

      # Some computation that creates intermediate `any_padded_graphs_tuple`s
      # after each message passing step.
      any_padded_graphs_tuple = any_padded_graphs_tuple.replace(
          edges=any_padded_graphs_tuple.edges * tf.cast(
              edges_mask, tf.float32)[:, None],
          nodes=any_padded_graphs_tuple.nodes * tf.cast(
              nodes_mask, tf.float32)[:, None],
          globals=any_padded_graphs_tuple.globals * tf.cast(
              graphs_mask, tf.float32)[:, None],
      )
      ```

  Returns:
    A `graphs.GraphsTuple` padded up to the required values.
  """
  padded_sizes = GraphsTupleSize(pad_nodes_to, pad_edges_to, pad_graphs_to)

  # The strategy goes as follows:
  # 0. Make sure our `graphs_tuple` is at least 1 node and 1 graph smaller than
  #    the padded sizes (this is required for step 1).
  # 1. Pad with one graph with at least one node, that contains all padding
  #    edges, and padding nodes, this will guaranteed preserved computation for
  #    graphs in the input `GraphsTuple`.
  # 2. Pad up to `pad_graphs_to` with graphs with no nodes and no edges.
  # 3. Set the shapes of the padded tensors to be statically known. Otherwise
  #    tensorflow shape inference mechanism is not smart enough to realize that
  #    at this stage tensors have statically known sizes.

  # Step 0.
  sufficient_space_assert = _assert_if_space_for_first_padding_graph(
      graphs_tuple, padded_sizes)
  with tf.control_dependencies([sufficient_space_assert]):
    padding_size = _get_required_padding_sizes(graphs_tuple, padded_sizes)

  # Step 1.
  first_padding_graph = _get_first_padding_graph(
      graphs_tuple, padding_size, experimental_unconnected_padding_edges)
  graphs_tuple_with_first_padding_graph = concat(
      [graphs_tuple, first_padding_graph], axis=0)

  # Step 2.
  remaining_padding_sizes = _get_required_padding_sizes(
      graphs_tuple_with_first_padding_graph, padded_sizes)
  padded_batch_kwargs = {}
  for field, tensor_dict in (
      graphs_tuple_with_first_padding_graph._asdict().items()):  # pylint: disable=protected-access
    field_pad_fn = functools.partial(
        _pad_tensor, padding_size=remaining_padding_sizes, field=field)
    padded_batch_kwargs[field] = tree.map_structure(field_pad_fn, tensor_dict)

  # Step 3.
  def _set_shape(tensor, padded_size):
    tensor_shape = tensor.get_shape().as_list()
    tensor.set_shape([padded_size] + tensor_shape[1:])
    return tensor
  for field, tensor_dict in padded_batch_kwargs.items():
    padded_size = _get_field_size_from_size_tuple(padded_sizes, field)
    set_shape_partial = functools.partial(_set_shape, padded_size=padded_size)
    tensor_dict = tree.map_structure(set_shape_partial, tensor_dict)
    padded_batch_kwargs[field] = tensor_dict
  return graphs.GraphsTuple(**padded_batch_kwargs)
Pads a `graphs.GraphsTuple` to fixed number of nodes, edges and graphs.

  The Graph Nets library treat zeros as valid parts of a graph.GraphsTuple, so
  special padding is required in order to preserve the computation. This
  method does so by adding a 'dummy' graph to the batch so that additional
  nodes/edges can't interfere with the valid graph.

  Args:
    graphs_tuple: `graphs.GraphsTuple` batch of graphs.
    pad_nodes_to: the size to pad node determined features to.
    pad_edges_to: the size to pad edge determined features to.
    pad_graphs_to: the size to pad graph determined features to.
    experimental_unconnected_padding_edges: Experimental feature to prevent nans
      in the padding graph. DISCLAIMER: This feature is extremly experimental,
      and setting it to `True` is not recommened unless strictly necessary, and
      understanding the implications.

      If `True`, the padding graph will have `senders` and `receivers` for
      the padding edges reference a node which does not exist (one beyond the
      size of `nodes`).

      This feature can be used to prevent any broadcasting/aggregation ops
      between edges and nodes for the padding graph. The reason is that the
      sum aggregations in the padding graph, which has a single node with a
      very large number of self-edges, sometimes lead to infs or nans,
      which may contaminate the gradients of the other valid graphs in the batch
      with nans (even if masked out of the loss: this is related to the
      `tf.where` issue.).

      This approach relies on several numerical hacks that do not work on CPU,
      but work on GPU and TPU (as covered by our tests):

      * `tf.gather` returns zeros when the index is beyond the boundaries. From
         https://www.tensorflow.org/api_docs/python/tf/gather
           "Note that on CPU, if an out of bound index is found, an error is
           returned. On GPU, if an out of bound index is found, a 0 is stored
           in the corresponding output value."
      * `tf.unsorted_segment_sum` drops values for negative indices. From
         https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum
           "If the given segment ID is negative, the value is dropped and
           will not be added to the sum of the segment."
        We have seen empirically that it also ignores values with indices equal
        or larger than `num_segments`. While this behavior is tested in our
        library, we cannot guarantee that it will work in the future for all
        unsorted_segment ops, so use at your own risk.

      This fixes the appearance of nans in the node-wise edge aggregation. The
      appearance of `nan`s is less likely in the global aggregation because in
      the worst case, the number of nodes/edges on the padding graph is not
      typically much larger than the number of nodes/edges in other graphs in
      the dataset.

      A less hacky approach (but more expensive, and requiring modifying model
      code) to prevent nan's appearing in the padding graph, is by masking out
      the graph features before they are aggregated, although for convenience
      we usually find that it is enough to do it after each message passing
      layer. E.g.:

      ```
      graphs_tuple_size = get_graphs_tuple_size(graphs_tuple)

      padded_graphs_tuple = pad_graphs_tuple(graphs_tuple, ...)

      graphs_mask = get_mask(graphs_tuple_size.num_graphs, pad_graphs_to)
      nodes_mask = get_mask(graphs_tuple_size.num_nodes, pad_nodes_to)
      edges_mask = get_mask(graphs_tuple_size.num_edges, pad_edges_to)

      # Some computation that creates intermediate `any_padded_graphs_tuple`s
      # after each message passing step.
      any_padded_graphs_tuple = any_padded_graphs_tuple.replace(
          edges=any_padded_graphs_tuple.edges * tf.cast(
              edges_mask, tf.float32)[:, None],
          nodes=any_padded_graphs_tuple.nodes * tf.cast(
              nodes_mask, tf.float32)[:, None],
          globals=any_padded_graphs_tuple.globals * tf.cast(
              graphs_mask, tf.float32)[:, None],
      )
      ```

  Returns:
    A `graphs.GraphsTuple` padded up to the required values.
  
graph_nets/utils_tf.py