Misc fixes for google import (#333)

This patch contains some misc fixes most related to type checking that
enable us to import gematria into google after the TF2 migration. Most
are pretty standard. We disable an attribute error on
run_continuous_evaluation mostly to silence it for now as tf_slim is not
really compatible with TF2 and needs to be rewritten.

Author: boomanaiden154

gematria/granite/python/gnn_model_base.py

@@ -90,9 +90,9 @@ class GraphNetworkLayer:
   num_iterations: Optional[int]
   layer_normalization: options.EnableFeature
   residual_connection: options.EnableFeature
-  edges_output_size: Sequence[int] = None
-  nodes_output_size: Sequence[int] = None
-  globals_output_size: Sequence[int] = None
+  edges_output_size: Sequence[int] | None = None
+  nodes_output_size: Sequence[int] | None = None
+  globals_output_size: Sequence[int] | None = None
 
 
 class GnnModelBase(model_base.ModelBase):
@@ -445,7 +445,9 @@ def _execute_graph_network(self, feed_dict) -> graph_nets.graphs.GraphsTuple:
     return graphs_tuple
 
   @abc.abstractmethod
-  def _execute_readout_network(self, graph_tuple) -> tf.Tensor:
+  def _execute_readout_network(
+      self, graph_tuple, feed_dict: model_base.FeedDict
+  ) -> tf.Tensor:
     """Creates a readout part of the network.
 
     Creates TensorFlow ops that take the output of the graph network and

gematria/granite/python/graph_builder_model_base.py

@@ -208,14 +208,18 @@ def _start_batch(self) -> None:
   # @Override
   def _make_batch_feed_dict(self) -> model_base.FeedDict:
     feed_dict = super()._make_batch_feed_dict()
-    feed_dict['instruction_node_mask'] = np.array(
+    feed_dict['instruction_node_mask'] = tf.constant(
         self._batch_graph_builder.instruction_node_mask, dtype=bool
     )
-    self._instruction_node_mask = feed_dict['instruction_node_mask']
-    feed_dict['instruction_annotations'] = (
+    self._instruction_node_mask = tf.constant(
+        feed_dict['instruction_node_mask']
+    )
+    feed_dict['instruction_annotations'] = tf.constant(
         self._batch_graph_builder.instruction_annotations
     )
-    self._instruction_annotations = feed_dict['instruction_annotations']
+    self._instruction_annotations = tf.constant(
+        feed_dict['instruction_annotations']
+    )
     return feed_dict
 
   # @Override

gematria/model/python/model_base.py

@@ -1113,8 +1113,10 @@ def run_continuous_evaluation(
         tf_slim.evaluation.StopAfterNEvalsHook(1),
         tf_slim.evaluation.SummaryAtEndHook(summary_dir, feed_dict=schedule),
         # Save the models with the best MAPE.
+        # We disable attribute error detection on self._loss because it is
+        # nullable and pytype expects there to be a check here.
         SaveBestCheckpoint(
-            error_tensor=self._loss.mean_absolute_percentage_error,
+            error_tensor=self._loss.mean_absolute_percentage_error,  # pytype: disable=attribute-error
             checkpoint_dir=os.path.join(summary_dir, 'best_models'),
             global_step=self.global_step,
         ),