More typing changes for google import (#335)

This patch contains a couple more typing changes that are needed to make
the google import work with the pytype checks. I have validated that
after this patch all of the tests run successfully internally.

Author: boomanaiden154

gematria/model/python/BUILD.bazel

@@ -104,6 +104,7 @@ gematria_py_library(
         "//gematria/basic_block/python:basic_block",
         "//gematria/basic_block/python:throughput",
         "//gematria/utils/python:timer",
+        "@graph_nets_repo//:graph_nets",
     ],
 )
 

gematria/model/python/model_base.py

@@ -45,12 +45,15 @@
 import scipy.stats
 import tensorflow as tf
 import tf_slim.evaluation
+import graph_nets
 
 # The type used for TensorFlow feed_dict objects. The type we use here is
 # simpler than what is actually accepted by TensorFlow, but the typing should be
 # sufficient for our use. Moreover, since TensorFlow and NumPy do not provide
 # type annotations, both the key and the value are reduced to typing.Any.
-FeedDict = MutableMapping[str, Union[np.ndarray, tf.Tensor]]
+FeedDict = MutableMapping[
+    str, Union[np.ndarray, tf.Tensor, graph_nets.graphs.GraphsTuple]
+]
 
 # A throughput value used as a placeholder in the expected output tensors for
 # masked expected outputs.
@@ -82,7 +85,7 @@ def __init__(
       self,
       error_tensor: tf.Tensor,
       checkpoint_dir: str,
-      global_step: int,
+      global_step: tf.Tensor,
       max_to_keep: int = 15,
   ):
     """Initializes the hook.
@@ -581,7 +584,7 @@ def _make_spearman_correlations(
       )
     return tf.concat(task_correlations, axis=0)
 
-  def _clip_if_not_none(self, grad: Optional[tf.Tensor]) -> tf.Tensor:
+  def _clip_if_not_none(self, grad: Optional[tf.Tensor]) -> tf.Tensor | None:
     if grad is None:
       return grad
     return tf.clip_by_norm(grad, self._grad_clip_norm)

gematria/model/python/training.py

@@ -70,17 +70,17 @@ class TrainingEpochStats:
       [len(this.percentile_ranks), num_tasks].
   """
 
-  epoch: int
+  epoch: tf.Tensor
   loss: float
   percentile_ranks: Sequence[int]
-  absolute_mse: np.ndarray
-  relative_mae: np.ndarray
-  relative_mse: np.ndarray
-  absolute_error_percentiles: np.ndarray
-  relative_error_percentiles: np.ndarray
-  absolute_delta_mse: Optional[np.ndarray] = None
-  absolute_delta_mae: Optional[np.ndarray] = None
-  absolute_delta_error_percentiles: Optional[np.ndarray] = None
+  absolute_mse: tf.Tensor
+  relative_mae: tf.Tensor
+  relative_mse: tf.Tensor
+  absolute_error_percentiles: tf.Tensor
+  relative_error_percentiles: tf.Tensor
+  absolute_delta_mse: Optional[tf.Tensor] = None
+  absolute_delta_mae: Optional[tf.Tensor] = None
+  absolute_delta_error_percentiles: Optional[tf.Tensor] = None
 
   def __post_init__(self) -> None:
     """Validate and finalize the initialization of TrainingEpochStats.