Add missing typing.Optional type annotations to function parameters.

PiperOrigin-RevId: 376238040

Author: rchen152

src/python/tensorflow_cloud/tuner/cloud_fit_client.py

@@ -45,12 +45,12 @@
 def cloud_fit(
     model: tf.keras.Model,
     remote_dir: Text,
-    region: Text = None,
-    project_id: Text = None,
-    image_uri: Text = None,
+    region: Optional[Text] = None,
+    project_id: Optional[Text] = None,
+    image_uri: Optional[Text] = None,
     distribution_strategy: Text = DEFAULT_DISTRIBUTION_STRATEGY,
-    job_spec: Dict[str, Any] = None,
-    job_id: Text = None,
+    job_spec: Optional[Dict[str, Any]] = None,
+    job_id: Optional[Text] = None,
     **fit_kwargs
 ) -> Text:
     """Executes in-memory Model and Dataset remotely on AI Platform.
@@ -209,7 +209,7 @@ def _serialize_assets(remote_dir: Text,
 def _default_job_spec(
     region: Text,
     image_uri: Text,
-    entry_point_args: Sequence[Text] = None,
+    entry_point_args: Optional[Sequence[Text]] = None,
 ) -> Dict[str, Any]:
     """Creates a basic job_spec for cloud AI Training.
 

src/python/tensorflow_cloud/tuner/tuner.py

@@ -93,10 +93,10 @@ def __init__(
         self,
         project_id: Text,
         region: Text,
-        objective: Union[Text, oracle_module.Objective] = None,
-        hyperparameters: hp_module.HyperParameters = None,
+        objective: Optional[Union[Text, oracle_module.Objective]] = None,
+        hyperparameters: Optional[hp_module.HyperParameters] = None,
         study_config: Optional[Dict[Text, Any]] = None,
-        max_trials: int = None,
+        max_trials: Optional[int] = None,
         study_id: Optional[Text] = None,
     ):
         """KerasTuner Oracle interface implemented with Vizier backend.
@@ -458,10 +458,10 @@ def __init__(
                                    tf.keras.Model]],
         project_id: Text,
         region: Text,
-        objective: Union[Text, oracle_module.Objective] = None,
-        hyperparameters: hp_module.HyperParameters = None,
+        objective: Optional[Union[Text, oracle_module.Objective]] = None,
+        hyperparameters: Optional[hp_module.HyperParameters] = None,
         study_config: Optional[Dict[Text, Any]] = None,
-        max_trials: int = None,
+        max_trials: Optional[int] = None,
         study_id: Optional[Text] = None,
         **kwargs):
         """Constructor.
@@ -516,10 +516,10 @@ def __init__(
         project_id: Text,
         region: Text,
         directory: Text,
-        objective: Union[Text, oracle_module.Objective] = None,
-        hyperparameters: hp_module.HyperParameters = None,
+        objective: Optional[Union[Text, oracle_module.Objective]] = None,
+        hyperparameters: Optional[hp_module.HyperParameters] = None,
         study_config: Optional[Dict[Text, Any]] = None,
-        max_trials: int = None,
+        max_trials: Optional[int] = None,
         study_id: Optional[Text] = None,
         container_uri: Optional[Text] = None,
         replica_config: Optional[machine_config.MachineConfig] = None,

src/python/tensorflow_cloud/tuner/utils.py

@@ -380,7 +380,7 @@ def convert_hyperparams_to_hparams(
 def format_objective(
     objective: Union[Text, oracle_module.Objective,
                      List[Union[Text, oracle_module.Objective]]],
-    direction: Text = None) -> List[oracle_module.Objective]:
+    direction: Optional[Text] = None) -> List[oracle_module.Objective]:
     """Formats objective to a list of oracle_module.Objective.
 
     Args:

src/python/tensorflow_cloud/tuner/vizier_client.py

@@ -39,7 +39,7 @@ def __init__(self,
                  service_client: discovery.Resource,
                  project_id: Text,
                  region: Text,
-                 study_id: Text = None):
+                 study_id: Optional[Text] = None):
         """Create an VizierClient object.
 
         Use this constructor when you know the study_id, and when the Study
@@ -206,10 +206,11 @@ def should_trial_stop(self, trial_id: Text) -> bool:
             return True
         return False
 
-    def complete_trial(self,
-                       trial_id: Text,
-                       trial_infeasible: bool,
-                       infeasibility_reason: Text = None)  -> Dict[Text, Any]:
+    def complete_trial(
+        self,
+        trial_id: Text,
+        trial_infeasible: bool,
+        infeasibility_reason: Optional[Text] = None)  -> Dict[Text, Any]:
         """Marks the trial as COMPLETED and sets the final measurement.
 
         Args:
@@ -289,7 +290,7 @@ def list_studies(self) -> List[Dict[Text, Any]]:
             raise
         return resp.get("studies", [])
 
-    def delete_study(self, study_name: Text = None) -> None:
+    def delete_study(self, study_name: Optional[Text] = None) -> None:
         """Deletes the study.
 
         Args:

src/python/tensorflow_cloud/tuner/vizier_client_interface.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """An abstract class for the client used in both OSS Vizier and Cloud AI Platform Optimizer Service."""
 import abc
-from typing import List, Mapping, Text, Union, Dict, Any
+from typing import Any, Dict, List, Mapping, Optional, Text, Union
 
 
 class VizierClientInterface(abc.ABC):
@@ -77,10 +77,11 @@ def should_trial_stop(self, trial_id: Text) -> bool:
     """
 
   @abc.abstractmethod
-  def complete_trial(self,
-                     trial_id: Text,
-                     trial_infeasible: bool,
-                     infeasibility_reason: Text = None) -> Dict[Text, Any]:
+  def complete_trial(
+      self,
+      trial_id: Text,
+      trial_infeasible: bool,
+      infeasibility_reason: Optional[Text] = None) -> Dict[Text, Any]:
     """Marks the trial as COMPLETED and sets the final measurement.
 
     Args:
@@ -110,7 +111,7 @@ def list_studies(self) -> List[Dict[Text, Any]]:
     """
 
   @abc.abstractmethod
-  def delete_study(self, study_name: Text = None) -> None:
+  def delete_study(self, study_name: Optional[Text] = None) -> None:
     """Deletes the study.
 
     Args: