Update pages/15_Graphs.py

pages/15_Graphs.py

@@ -5,6 +5,10 @@ from tensorflow_gnn.models import mt_albis
 import networkx as nx
 import matplotlib.pyplot as plt
 
+# Set environment variable for legacy Keras
+import os
+os.environ['TF_USE_LEGACY_KERAS'] = '1'
+
 # Define the model function
 def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     graph = inputs = tf.keras.Input(type_spec=graph_tensor_spec)
@@ -17,10 +21,15 @@ def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     # Message passing layers
     for _ in range(2):
         graph = mt_albis.MtAlbisGraphUpdate(
-            units=128, message_dim=64,
-            attention_type="none", simple_conv_reduce_type="mean",
-            normalization_type="layer", next_state_type="residual",
-            state_dropout_rate=0.2, l2_regularization=1e-5,
+            units=128,
+            message_dim=64,
+            receiver_tag=tfgnn.NODES,  # Add this line
+            attention_type="none",
+            simple_conv_reduce_type="mean",
+            normalization_type="layer",
+            next_state_type="residual",
+            state_dropout_rate=0.2,
+            l2_regularization=1e-5,
         )(graph)
     
     return tf.keras.Model(inputs, graph)
@@ -47,8 +56,8 @@ def create_sample_graph():
             "cites": tfgnn.EdgeSet.from_fields(
                 sizes=[num_edges],
                 adjacency=tfgnn.Adjacency.from_indices(
-                    source=("papers", tf.cast(list(e[0] for e in graph.edges()), tf.int32)),
-                    target=("papers", tf.cast(list(e[1] for e in graph.edges()), tf.int32))
+                    source=("papers", tf.cast([e[0] for e in graph.edges()], tf.int32)),
+                    target=("papers", tf.cast([e[1] for e in graph.edges()], tf.int32))
                 ),
                 features={"features": edge_features}
             )