Why tfgnn.gather_first_node() can't get values for mutiple node sets in one time?

[SidneyLann]

    def readout_features(self, graphTensor: tfgnn.GraphTensor, feature_name: str) -> tf.Tensor:
    """Extract features using dynamic root node type from context"""
    # 1. Get root node type from context feature
    root_node_type = graphTensor.context["root_node_type"]  # Shape [batch_size, 1]
    #root_node_type = tf.squeeze(root_node_type, axis=1)      # Shape [batch_size]
    
    batch_size = tf.shape(root_node_type)[0]
    features = tf.TensorArray(tf.float32, size=batch_size)
    
    for i in tf.range(batch_size):
            current_type = root_node_type[i]
            component = tfgnn.get_component(graph_tensor, i)
            feature = tfgnn.gather_first_node(
                    component,
                    node_set_name=current_type,
                    feature_name=tfgnn.HIDDEN_STATE
            )
            features = features.write(i, feature)
    
    return features.stack()

It is very inconvenient to get embeddings for first nodes mutiple node sets in one time.

[arnoegw]

It sounds as if tfgnn.structured_readout() could help here. Please see the respective sections of the Data Preparation and GNN Modeling guides as well as the docstrings of tfgnn.structured_readout() function and the equivalent tfgnn.keras.layers.StructuredReadout layer.

[SidneyLann]

@arnoegw this function also only get for one NodeSetName but not NodeSetName collection.
tfgnn.keras.layers.StructuredReadout(
key: Optional[str] = None,
*,
feature_name: str = const.HIDDEN_STATE,
readout_node_set: NodeSetName = '_readout',
validate: bool = True,
**kwargs
)

I now use seperate model per node set for my 5 node sets gnn. However, there are 4 logs WARNING:tensorflow:Gradients do not exist for variables, it seem that other 4 node sets missing in specific model.

[arnoegw]

The trick is to set up a readout structure (the "_readout" node set and specifically named edge sets into it) to encode from where the readout is meant to happen, perhaps differently for each input graph.

[SidneyLann]

The root nodes must be all the same type for an input graph? if no, it can't read all root states in one times?