[𝘀𝗽𝗿] changes to main this commit is based on (#337)

We cannot make the dataclass frozen when doing this, but this enables TF
to automatically find trainable variables within GraphNetworkLayer
objects which means we can get rid of the somewhat hacky
_get_trainable_variables function that subclasses were supposed to
override. This successfully trains models that would otherwise fail to
converge if none of the graph layers were trainable.

This closes #323.

Author: boomanaiden154

gematria/granite/python/gnn_model_base.py

@@ -30,8 +30,8 @@
 import tf_keras
 
 
-@dataclasses.dataclass(frozen=True)
-class GraphNetworkLayer:
+@dataclasses.dataclass
+class GraphNetworkLayer(tf.Module):
   """Specifies one segment of the pipeline of the graph network.
 
   Each segment consists of a graph network module, i.e. a Sonnet module that
@@ -290,12 +290,6 @@ def initialize(self):
               tf_keras.layers.LayerNormalization(name=globals_layer_norm_name)
           )
 
-  def _get_trainable_variables(self):
-    trainable_variables = list(super()._get_trainable_variables())
-    for layer in self._graph_network:
-      trainable_variables.extend(layer.module.trainable_variables)
-    return trainable_variables
-
   # @Override
   def _forward(self, feed_dict):
     graph_tuple_outputs = self._execute_graph_network(feed_dict)

gematria/model/python/model_base.py

@@ -1331,9 +1331,6 @@ def compute_loss_tensor(self, schedule: FeedDict):
         )
     )
 
-  def _get_trainable_variables(self):
-    return self.trainable_variables
-
   def train_batch(
       self,
       schedule: FeedDict,
@@ -1367,11 +1364,10 @@ def train_batch(
               for variable in self._variable_groups.get(variable_group)
           )
 
-        trainable_variables = self._get_trainable_variables()
         variables = (
             [variable.deref() for variable in requested_variables]
             if requested_variables
-            else trainable_variables
+            else self.trainable_variables
         )
 
         grads = tape.gradient(loss_tensor, variables)