operations between two GraphTuple

[TianrenWang]

Let's say I have two GraphTuples with the same number of nodes and the same connectivity, is there a function that returns the product/sum/difference between the two? The specific use case for me is I want a way to only perform graph computation on the graph states rather than on the graph features, but I also want the graph features to guide the computation performed on the graph states by using the product/sum/difference operations.

[zafarali]

Thanks for your question!

One option could be to use tree.map_structure:

import tree
from graph_nets.graphs import GraphsTuple
import numpy as np

a = GraphsTuple(nodes=np.ones((5, 1)), n_node=[5], senders=None, edges=None, receivers=None, globals=None, n_edge=[0])
b = GraphsTuple(nodes=2 * np.ones((5, 1)), n_node=[5], senders=None, edges=None, receivers=None, globals=None, n_edge=[0])

def negative(graph_1, graph_2):
  if graph_1 is None:
    return None
  return graph_1 - graph_2

tree.map_structure(negative, a, b)

which outputs:

GraphsTuple(nodes=array([[-1.],
       [-1.],
       [-1.],
       [-1.],
       [-1.]]), edges=None, receivers=None, senders=None, globals=None, n_node=[0], n_edge=[0])

You will need to be careful not to apply negative on the senders and receivers.

Another option could be: a.replace(nodes=a.nodes - b.nodes) which should give you the same thing.