Override gnn_model_base trainable_variables

This patch provides a custom implementation of trainable_variables in
gnn_model_base. Theoretically this should have been made unnecessary by
\#337, but the interanl version of Tensorflow refuses to recurse into
the modules inside of the GraphNetworkLayer classes. This patch fixes
that by just returning the values regardless.

Reviewers: orodley, virajbshah, ondrasej

Reviewed By: ondrasej

Pull Request: #341

Author: boomanaiden154

gematria/granite/python/gnn_model_base.py

@@ -210,6 +210,14 @@ def __init__(
     self._graph_module_residual_connections = graph_module_residual_connections
     self._graph_module_layer_normalization = graph_module_layer_normalization
 
+  @property
+  def trainable_variables(self):
+    trainable_vars = set(var.ref() for var in super().trainable_variables)
+    for layer in self._graph_network:
+      layer_vars = [var.ref() for var in layer.module.trainable_variables]
+      trainable_vars.update(layer_vars)
+    return tuple(var.deref() for var in trainable_vars)
+
   def initialize(self):
     super().initialize()
     self._graph_network = self._create_graph_network_modules()