How can I transform 6 separate graph tensor into a single tf.data.Dataset and batch it?

[cherryCookies]

To whom it may concern,

Before I get into the issue, please bear with me if I get lost since I am a novice in the field of graph neural network, and I am on my own. I am struggling with transforming 6 separate graph tensor into a single dataset of a 6 batches. The graph schema for each heterogenous graph is identical, and they consist of one node sets and two edge sets. (Sorry for not being able to share the exact schema because it is a part of a research)

Now I would like to transform 6 graph tensor into a dataset with 6 batches. So, I have done following code to merge graph tensors into one single Datasets, so that I can perform Datasets.batch()

Graph tensor into tf.data.Dataset

dataset_l = [] for i in range(6): dataset = eachGraphTensor.map(lambda serialized:tfgnn.parse_single_example(input_spec, serialized)) dataset_l.append(dataset)

Merging tf.data.Dataset

dataset_l = [6 graph tensor in a list] mergedDataset = dataset_l[0] for i in range(1,len(dataset_l)): mergedDataset.concatenate(dataset_l[i])

Batching

dataset = mergedDataset.batch(1)

Is this acceptable way to merge individual graph tensor into a single dataset and batching it?

Kind Regards Charyeong Heo from Seoul, Korea