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

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

Start with the shown tfgnn.parse_single_example, Dataset.concatenate, and Dataset.batch calls, and verify their behavior using six equivalent heterogeneous graph datasets. No repository files or tests are named; done means determining whether the resulting dataset represents six graph elements and batches them as intended, including the effect of the reported construction.