This is my Graph tensor.

featureA_dim = 768
featureB_dim = 61
featureC_dim = 1
def create_graph_tensor(node_df,edge_df_tmp,genomic_edge_df_tmp):
    graph_tensor = tfgnn.GraphTensor.from_pieces(
        node_sets = {"repeats":
                        tfgnn.NodeSet.from_fields(sizes = [-1],
                        features ={
                            'featureA':np.array(featureA_data, dtype='float32').reshape(-1,featureA_dim),
                            'featureB':np.array(featureB_data, dtype='float32').reshape(-1,featureB_dim),
                            'featureC':np.array(featureC_data, dtype='float32').reshape(-1,featureC_dim)
})},
        edge_sets = {"order":
                        tfgnn.EdgeSet.from_fields(sizes = [len(genomic_edge_df_tmp)],
                        features = {},
                        adjacency = tfgnn.Adjacency.from_indices(
                            source = ("repeats", np.array(genomic_edge_df_tmp['previous_node_id'], dtype='int32').reshape(-1,1)),
                            target = ("repeats", np.array(genomic_edge_df_tmp['following_node_id'], dtype='int32').reshape(-1,1))                                     )),
                    "possibility": tfgnn.EdgeSet.from_fields(sizes = [len(edge_df_tmp)],
                    features = {'label': np.array(edge_df_tmp['label'], dtype='int32').reshape(-1,1)},
                    adjacency = tfgnn.Adjacency.from_indices(
                            source = ("repeats", np.array(edge_df_tmp['previous_node_id'], dtype='int32')),
                            target = ("repeats", np.array(edge_df_tmp['following_node_id'], dtype='int32'))))})
    return graph_tensor

For the above graph tensor, I have processed the edge to separate the label.

def edge_merge_encode(graph):
    edge_features = graph.edge_sets['possibility'].get_features_dict()
    label = edge_features.pop('label')
    new_graph = graph.replace_features(edge_sets={'possibility':edge_features})
    return new_graph, label

def create_dataset(graph,function):
    dataset = tf.data.Dataset.from_tensors(graph)
    return dataset.map(function)

#INPUT:DO AWAY LABELS(Edge Datasets)
full_edge_dataset = create_dataset(full_tensor,edge_merge_encode)
train_edge_dataset = create_dataset(train_tensor,edge_merge_encode)

The train_edge_dataset has GraphTensorSpec as follows.

<_MapDataset element_spec=(GraphTensorSpec({'context': ContextSpec({'features': {}, 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None), 'node_sets': {'repeats': NodeSetSpec({'features': {'featureA': TensorSpec(shape=(942, 768), dtype=tf.float32, name=None), 'featureB': TensorSpec(shape=(942, 61), dtype=tf.float32, name=None), 'featureC': TensorSpec(shape=(942, 1), dtype=tf.float32, name=None)}, 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None)}, 'edge_sets': {'order': EdgeSetSpec({'features': {}, 'adjacency': AdjacencySpec({'#index.0': TensorSpec(shape=(939, 1), dtype=tf.int32, name=None), '#index.1': TensorSpec(shape=(939, 1), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, {'#index.0': 'repeats', '#index.1': 'repeats'}), 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None), 'possibility': EdgeSetSpec({'features': {}, 'adjacency': AdjacencySpec({'#index.0': TensorSpec(shape=(282,), dtype=tf.int32, name=None), '#index.1': TensorSpec(shape=(282,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, {'#index.0': 'repeats', '#index.1': 'repeats'}), 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None)}}, TensorShape([]), tf.int32, tf.int64, None), TensorSpec(shape=(282, 1), dtype=tf.int32, name=None))>

I have trained train_edge_dataset with below graph-updating layer architecture.

`

def set_initial_node_state(node_set, node_set_name):
    features = [tf.keras.layers.Dense(64,activation="relu")(node_set['featureA']),
                tf.keras.layers.Dense(32,activation="relu")(node_set['featureB']),
                tf.keras.layers.Dense(2,activation="relu")(node_set['featureC']) ]
    return tf.keras.layers.Concatenate()(features)

def set_initial_edge_state(edge_set, edge_set_name):
    if edge_set_name == "order":
        features = tfgnn.keras.layers.MakeEmptyFeature()(edge_set)
        return features
    elif edge_set_name == "possibility":
        features = tfgnn.keras.layers.MakeEmptyFeature()(edge_set)
        return features

graph_spec,_ = train_edge_dataset.element_spec
input_graph = tf.keras.layers.Input(type_spec=graph_spec)

graph = tfgnn.keras.layers.MapFeatures(
    node_sets_fn=set_initial_node_state,
    edge_sets_fn=set_initial_edge_state
    )(input_graph)

#MODEL:SPCIFYING LAYERS
graph_updates = 3
for i in range(graph_updates):
    graph = tfgnn.keras.layers.GraphUpdate(
        edge_sets = {
                'possibility': tfgnn.keras.layers.EdgeSetUpdate(
                    next_state = tfgnn.keras.layers.NextStateFromConcat( tf.keras.layers.Dense(8))),
        },
        node_sets = {
                'repeats': tfgnn.keras.layers.NodeSetUpdate(
                    {
                        "order": tfgnn.keras.layers.SimpleConv(
                            tf.keras.layers.Dense(32, "relu"),
                            "mean",
                            receiver_tag=tfgnn.TARGET),
                        "possibility": tfgnn.keras.layers.SimpleConv(
                            tf.keras.layers.Dense(32, "relu"),
                            "mean",
                            receiver_tag=tfgnn.TARGET)
                    },
tfgnn.keras.layers.NextStateFromConcat(tf.keras.layers.Dense(64)))})(graph)


#OUTPUT:FINAL OUTPUT LAYER
logits = tf.keras.layers.Dense(8,activation='relu')(graph.edge_sets['possibility'][tfgnn.HIDDEN_STATE])
logits = tf.keras.layers.Dense(1,activation='sigmoid')(logits)

The model was compiled as follows.

edge_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.00001),
    loss=['binary_crossentropy'],
    metrics=['Accuracy'])

es = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',
    mode='min',
    verbose=1,
    patience=1000,
    restore_best_weights=True)

history = edge_model.fit(train_edge_dataset,
                            epochs=10,
                            validation_data=val_edge_dataset,
                            callbacks=[es],
                            verbose=1)

Training with train_edge_dataset went well, but training with full_edge_dataset is problematic.
The GraphTensorSpec for full_edge_dataset is as follows.

<_MapDataset element_spec=(GraphTensorSpec({'context': ContextSpec({'features': {}, 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None), 'node_sets': {'repeats': NodeSetSpec({'features': {'featureA': TensorSpec(shape=(942, 768), dtype=tf.float32, name=None), 'featureB': TensorSpec(shape=(942, 61), dtype=tf.float32, name=None), 'featureC': TensorSpec(shape=(942, 1), dtype=tf.float32, name=None)}, 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None)}, 'edge_sets': {'order': EdgeSetSpec({'features': {}, 'adjacency': AdjacencySpec({'#index.0': TensorSpec(shape=(939, 1), dtype=tf.int32, name=None), '#index.1': TensorSpec(shape=(939, 1), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, {'#index.0': 'repeats', '#index.1': 'repeats'}), 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None), 'possibility': EdgeSetSpec({'features': {}, 'adjacency': AdjacencySpec({'#index.0': TensorSpec(shape=(376,), dtype=tf.int32, name=None), '#index.1': TensorSpec(shape=(376,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, {'#index.0': 'repeats', '#index.1': 'repeats'}), 'sizes': TensorSpec(shape=(1,), dtype=tf.int32, name=None)}, TensorShape([]), tf.int32, tf.int64, None)}}, TensorShape([]), tf.int32, tf.int64, None), TensorSpec(shape=(376, 1), dtype=tf.int32, name=None))>

I have checked the label and logit shape, and I see no problem. The logits for full_edge_dataset could not be retrieved because the full_edge_dataset could not go through the layers. (the task that I want is link prediction case)

Shape of label for full_edge_dataset tensor during model fitting: (376, 1)
Shape of label for train_edge_dataset tensor during model fitting: (282, 1)
Shape of logits for train_edge_dataset after final layer: (282, 1)

The exact error message is as follows.

Traceback (most recent call last):
File "gnn_final.py", line 282, in <module>edge_model.fit(full_edge_dataset,
File "~/.local/lib/python3.8/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handlerraise e.with_traceback(filtered_tb) from None
File "~/.local/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 53, in quick_executetensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

Detected at node 'IteratorGetNext' defined at (most recent call last):
File "gnn_final.py", line 282, in <module>edge_model.fit(full_edge_dataset,
File "~/.local/lib/python3.8/site-packages/keras/src/utils/traceback_utils.py", line 65, in error_handler
return fn(*args, **kwargs)
File "~/.local/lib/python3.8/site-packages/keras/src/engine/training.py", line 1742, in fit
tmp_logs = self.train_function(iterator)
File "~/.local/lib/python3.8/site-packages/keras/src/engine/training.py", line 1338, in train_function
return step_function(self, iterator)
File "~/.local/lib/python3.8/site-packages/keras/src/engine/training.py", line 1321, in step_function
data = next(iterator)
Node: 'IteratorGetNext'Detected at node 'IteratorGetNext' defined at (most recent call last):
File "gnn_final.py", line 282, in <module>edge_model.fit(full_edge_dataset,File "~/.local/lib/python3.8/site-packages/keras/src/utils/traceback_utils.py", line 65, in error_handler
return fn(*args, **kwargs)
File "~/.local/lib/python3.8/site-packages/keras/src/engine/training.py", line 1742, in fit
tmp_logs = self.train_function(iterator)
File "~/.local/lib/python3.8/site-packages/keras/src/engine/training.py", line 1338, in train_function
return step_function(self, iterator)File "~/.local/lib/python3.8/site-packages/keras/src/engine/training.py", line 1321, in step_function
data = next(iterator)
Node: 'IteratorGetNext'
2 root error(s) found.
(0) INVALID_ARGUMENT:  Data type mismatch at component 4: expected variant but got int32.
[[{{node IteratorGetNext}}]]
[[IteratorGetNext/_4]]
(1) INVALID_ARGUMENT:  Data type mismatch at component 4: expected variant but got int32.[[{{node IteratorGetNext}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_12592]