Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -155,28 +155,23 @@ def set_initial_node_states(node_set: tfgnn.NodeSet, node_set_name: str):
         return node_set["empty_state"]
     raise KeyError(f"Unexpected node_set_name='{node_set_name}'")
 
-def gnn_model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
+def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
     graph = tfgnn.keras.layers.MapFeatures(node_sets_fn=set_initial_node_states)(inputs)
-    for i in range(num_graph_updates):
+    
+    for _ in range(num_graph_updates):
         graph = mt_albis.MtAlbisGraphUpdate(
             units=node_state_dim,
             message_dim=message_dim,
-            receiver_tag=tfgnn.SOURCE,
-            node_set_names=None if i < num_graph_updates-1 else ["paper"],
+            attention_type="none",
             simple_conv_reduce_type="mean|sum",
-            state_dropout_rate=state_dropout_rate,
-            l2_regularization=l2_regularization,
             normalization_type="layer",
             next_state_type="residual",
+            state_dropout_rate=state_dropout_rate,
+            l2_regularization=l2_regularization
         )(graph)
-    return tf.keras.Model(inputs, graph)
-
-def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
-    gnn_model = gnn_model_fn(graph_tensor_spec)
-    inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
-    graph = gnn_model(inputs)
-    paper_state = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="hidden_state")(graph)
+    
+    paper_state = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="state")(graph)
     paper_state = tf.keras.layers.Dense(349, activation="softmax")(paper_state)
     return tf.keras.Model(inputs, paper_state)
 
@@ -215,7 +210,7 @@ initial_learning_rate = 0.001
 steps_per_epoch = NUM_TRAINING_SAMPLES // global_batch_size
 validation_steps = NUM_VALIDATION_SAMPLES // global_batch_size
 learning_rate = tf.keras.optimizers.schedules.CosineDecay(
-    initial_learning_rate, steps_per_epoch*epochs)
+    initial_learning_rate, steps_per_epoch * epochs)
 optimizer_fn = functools.partial(tf.keras.optimizers.Adam, learning_rate=learning_rate)
 
 # Define trainer