class WrappedGCN(tf.keras.layers.Layer):
    def __init__(self, features, *args, **kwargs):
        super(WrappedGCN, self).__init__()
        self.features = features
        self.layer = GCNConv(features, *args, **kwargs)
    
    def hook(self, graph):
        features = graph.features.to_tensor()
        features = tf.reshape(features, (1, -1, features.shape[1]))
        adj = tf.cast(graph.adjacency.to_tensor(), tf.float32)
        adj = tf.reshape(adj, (1, adj.shape[0], adj.shape[1]))
        return tf.RaggedTensor.from_tensor(tf.squeeze(self.layer([features, adj])))

    def __call__(self, graph):
        if isinstance(graph, TensorGraph):
            features = tf.map_fn(self.hook, graph, tf.RaggedTensorSpec(
                shape=(None, self.features), dtype=tf.float32))
            return TensorGraph(
                features=features,
                adjacency=graph.adjacency)
        return self.layer(graph)

x0 = TensorGraph(adjacency=tf.ragged.stack(adjs),
                    features=tf.ragged.stack(features))

x1 = WrappedGCN(6)(x0)
x2 = WrappedGCN(6)(x1)

class TensorGraph(tf.experimental.BatchableExtensionType):
    """A collection of nodes with associated feature vectors."""
    features: tf.RaggedTensor
    adjacency: tf.RaggedTensor
    
    # TODO: Validation functions etc...