fix pytype

google · Aug 12, 2022 · 113c71e · 113c71e
1 parent 468adc3
commit 113c71e
Showing 1 changed file with 12 additions and 10 deletions.
diff --git a/compiler_opt/rl/compilation_runner.py b/compiler_opt/rl/compilation_runner.py
@@ -230,11 +230,12 @@ class CompilationRunnerStub(metaclass=abc.ABCMeta):
   """The interface of a stub to CompilationRunner, for type checkers."""
 
   @abc.abstractmethod
-  def collect_results(self,
-                      module_spec: corpus.ModuleSpec,
-                      tf_policy_path: str,
-                      collect_default_result: bool,
-                      reward_only: bool = False) -> Tuple[Dict, Dict]:
+  def collect_results(
+      self,
+      module_spec: corpus.ModuleSpec,
+      tf_policy_path: str,
+      collect_default_result: bool,
+      reward_only: bool = False) -> Tuple[Optional[Dict], Optional[Dict]]:
     raise NotImplementedError()
 
   @abc.abstractmethod
@@ -290,11 +291,12 @@ def get_rewards(result: Dict) -> List[float]:
       return []
     return [v[1] for v in result.values()]
 
-  def collect_results(self,
-                      module_spec: corpus.ModuleSpec,
-                      tf_policy_path: str,
-                      collect_default_result: bool,
-                      reward_only: bool = False) -> Tuple[Dict, Dict]:
+  def collect_results(
+      self,
+      module_spec: corpus.ModuleSpec,
+      tf_policy_path: str,
+      collect_default_result: bool,
+      reward_only: bool = False) -> Tuple[Optional[Dict], Optional[Dict]]:
     """Collect data for the given IR file and policy.
 
     Args: