Using TFRecord to store graphs

TensorFlow General Discussion
1

Hello,

I’m building a model using the Spektral library. My pipeline that produces datasets flattens the graph features, adjacency matrix, and edge features.

The structure looks like this:

    parsed_features = tf.io.parse_single_example(
        serialized=proto,
        features=dict(
            node_features=tf.io.VarLenFeature(dtype=tf.int64),
            node_features_shape=tf.io.VarLenFeature(dtype=tf.int64),
            adjacency_matrix=tf.io.VarLenFeature(dtype=tf.int64),
            adjacency_matrix_shape=tf.io.VarLenFeature(dtype=tf.int64),
            edge_features=tf.io.VarLenFeature(dtype=tf.int64),
            edge_features_shape=tf.io.VarLenFeature(dtype=tf.int64),
            label=tf.io.VarLenFeature(dtype=tf.int64),
        ),
    )

Each Example stores data of different lengths. To restore the original format, I decided to store shapes alongside the flattened data. What would be the most idiomatic way to bring node_features, adjacency_matrix and edge_features back to their original shapes?

I tried a couple of things, but dealing with SymbolicTensors makes it tricky.