code: def convert_to_graphs_tuple(
      graph: Dict[str, tf.Tensor],
      add_virtual_node: bool,
      add_undirected_edges: bool,
      add_self_loops: bool,
  ) -> jraph.GraphsTuple:
      """Converts a dictionary of tf.Tensors to a GraphsTuple."""
      num_nodes = tf.squeeze(graph["num_nodes"])
      num_edges = tf.squeeze(graph["num_edges"])
      nodes = graph["node_feat"]
      edges = graph["edge_feat"]
      edge_feature_dim = edges.shape[-1]
      labels = graph["labels"]
      senders = graph["edge_index"][:, 0]
      receivers = graph["edge_index"][:, 1]
  
      # Add a virtual node connected to all other nodes.
      # The feature vectors for the virtual node
      # and the new edges are set to all zeros.
      if add_virtual_node:
          nodes = tf.concat(
              [nodes, tf.zeros_like(nodes[0, None])], axis=0
          )
          senders = tf.concat([senders, tf.range(num_nodes)], axis=0)
          receivers = tf.concat(
              [receivers, tf.fill((num_nodes,), num_nodes + 1)], axis=0
          )
          edges = tf.concat(
              [edges, tf.zeros((num_nodes, edge_feature_dim))], axis=0
          )
          num_edges += num_nodes
          num_nodes += 1
  
      # Make edges undirected, by adding edges with senders and receivers flipped.
      # The feature vector for the flipped edge is the same as the original edge.
      if add_undirected_edges:
          new_senders = tf.concat([senders, receivers], axis=0)
          new_receivers = tf.concat([receivers, senders], axis=0)
          edges = tf.concat([edges, edges], axis=0)
          senders, receivers = new_senders, new_receivers
          num_edges *= 2
  
      # Add self-loops for each node.
      # The feature vectors for the self-loops are set to all zeros.
      if add_self_loops:
          senders = tf.concat([senders, tf.range(num_nodes)], axis=0)
          receivers = tf.concat(
              [receivers, tf.range(num_nodes)], axis=0
          )
          edges = tf.concat(
              [edges, tf.zeros((num_nodes, edge_feature_dim))], axis=0
          )
          num_edges += num_nodes
  
      return jraph.GraphsTuple(
          n_node=tf.expand_dims(num_nodes, 0),
          n_edge=tf.expand_dims(num_edges, 0),
          nodes=nodes,
          edges=edges,
          senders=senders,
          receivers=receivers,
          globals=tf.expand_dims(labels, axis=0),
      )
  
file_path: flax/examples/ogbg_molpcba/input_pipeline.py
imports: import functools
  from typing import Dict, NamedTuple
  import jraph
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: convert_to_graphs_tuple(graph, add_virtual_node, add_undirected_edges, add__loops) and returns: jraph