Training on batches of GraphsTuples?

Let's say I want to train an LSTM or transformer on sequences of graphs using Sonnet2/TF2:

I want to represent the graphs in each sequence as one ```GraphsTuple```, which means my batches are essentially an iterable of ```GraphsTuple```s, each with a variable number of graphs. This is great until it's time to get the input signature and compile the update step. It's unclear to me how to define the tensorspec for this type of input. Is my best route to subclass ```collections.namedtuple()``` similar to how you define a ```GraphsTuple```, or can you suggest a more elegant solution?

Thanks