Made api's compatible with python3.7+

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Keras-based TransformerEncoder block layer."""
-from typing import Any, Optional, Sequence
+from typing import Any, Optional, Sequence, Union
 from absl import logging
 import tensorflow as tf, tf_keras
 
@@ -28,9 +28,9 @@ class RMSNorm(tf_keras.layers.Layer):
 
   def __init__(
       self,
-      axis: int | Sequence[int] = -1,
+      axis: Union[int, Sequence[int]] = -1,
       epsilon: float = 1e-6,
-      **kwargs,
+      **kwargs
   ):
     """Initializes RMSNorm.
 
@@ -43,7 +43,8 @@ def __init__(
     self.axis = [axis] if isinstance(axis, int) else axis
     self.epsilon = epsilon
 
-  def build(self, input_shape: tf.TensorShape | Sequence[int | None]):
+  def build(self, 
+            input_shape: Union[tf.TensorShape, Sequence[Optional[int]]]):
     input_shape = tf.TensorShape(input_shape)
     scale_shape = [1] * input_shape.rank
     for dim in self.axis:

official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py

@@ -15,6 +15,7 @@
 """Defines base abstract uplift network layers."""
 
 import abc
+from typing import Union, Optional
 
 import tensorflow as tf, tf_keras
 
@@ -33,7 +34,7 @@ class BaseTwoTowerUpliftNetwork(tf_keras.layers.Layer, metaclass=abc.ABCMeta):
   def call(
       self,
       inputs: types.DictOfTensors,
-      training: bool | None = None,
-      mask: tf.Tensor | None = None,
+      training: Optional[bool] = None,
+      mask: Optional[tf.Tensor] = None,
   ) -> types.TwoTowerTrainingOutputs:
     raise NotImplementedError()

official/recommendation/uplift/metrics/label_mean.py

@@ -14,6 +14,8 @@
 
 """Keras metric for computing the label mean sliced by treatment group."""
 
+from typing import Union, Optional
+
 import tensorflow as tf, tf_keras
 
 from official.recommendation.uplift import types
@@ -71,7 +73,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Optional[tf.Tensor] = None,
   ):
     """Updates the overall, control and treatment label means.
 

official/recommendation/uplift/metrics/label_variance.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Keras metric for computing the label variance sliced by treatment group."""
+from typing import Optional
 
 import tensorflow as tf, tf_keras
 
@@ -72,7 +73,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Optional[tf.Tensor] = None,
   ):
     """Updates the overall, control and treatment label variances.
 

official/recommendation/uplift/metrics/metric_configs.py

@@ -16,7 +16,7 @@
 
 from collections.abc import Mapping
 import dataclasses
-from typing import Any
+from typing import Any, Optional
 
 from official.core.config_definitions import base_config
 
@@ -33,9 +33,9 @@ class SlicedMetricConfig(base_config.Config):
       values to slice on.
   """
 
-  slicing_feature: str | None = None
-  slicing_spec: Mapping[str, int] | None = None
-  slicing_feature_dtype: str | None = None
+  slicing_feature: Optional[str] = None
+  slicing_spec: Optional[Mapping[str, int]] = None
+  slicing_feature_dtype: Optional[str] = None
 
   def __post_init__(
       self, default_params: dict[str, Any], restrictions: list[str]

official/recommendation/uplift/metrics/sliced_metric.py

@@ -15,6 +15,7 @@
 """Keras metric for reporting metrics sliced by a feature."""
 
 import copy
+from typing import Union, Optional
 
 import tensorflow as tf, tf_keras
 
@@ -66,9 +67,9 @@ class SlicedMetric(tf_keras.metrics.Metric):
   def __init__(
       self,
       metric: tf_keras.metrics.Metric,
-      slicing_spec: dict[str, str] | dict[str, int],
-      slicing_feature_dtype: tf.DType | None = None,
-      name: str | None = None,
+      slicing_spec: Union[dict[str, str], dict[str, int]],
+      slicing_feature_dtype: Optional[tf.DType] = None,
+      name: Optional[str] = None,
   ):
     """Initializes the instance.
 
@@ -123,7 +124,7 @@ def __init__(
   def update_state(
       self,
       *args: tf.Tensor,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Optional[tf.Tensor] = None,
       slicing_feature: tf.Tensor,
       **kwargs,
   ):

official/recommendation/uplift/metrics/treatment_fraction.py

@@ -18,6 +18,7 @@
 
 from official.recommendation.uplift import types
 
+from typing import Optional
 
 @tf_keras.utils.register_keras_serializable(package="Uplift")
 class TreatmentFraction(tf_keras.metrics.Metric):
@@ -57,7 +58,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Optional[tf.Tensor] = None,
   ) -> None:
     """Updates the treatment fraction.
 

official/recommendation/uplift/metrics/uplift_mean.py

@@ -14,6 +14,8 @@
 
 """Keras metric for computing the mean uplift sliced by treatment group."""
 
+from typing import Optional
+
 import tensorflow as tf, tf_keras
 
 from official.recommendation.uplift import types
@@ -68,7 +70,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Optional[tf.Tensor] = None,
   ) -> None:
     """Updates the overall, control and treatment uplift means.
 

official/recommendation/uplift/types.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 """Defines types used by the keras uplift modeling library."""
+from typing import Union
 
 import tensorflow as tf, tf_keras
 
-TensorType = tf.Tensor | tf.SparseTensor | tf.RaggedTensor
-
+TensorType = Union[tf.Tensor, tf.SparseTensor, tf.RaggedTensor]
 ListOfTensors = list[TensorType]
 TupleOfTensors = tuple[TensorType, ...]
 DictOfTensors = dict[str, TensorType]
 
-CollectionOfTensors = ListOfTensors | TupleOfTensors | DictOfTensors
+CollectionOfTensors = Union[ListOfTensors, TupleOfTensors, DictOfTensors]
 
 
 class TwoTowerNetworkOutputs(tf.experimental.ExtensionType):

official/vision/configs/retinanet.py

@@ -53,7 +53,7 @@ class Parser(hyperparams.Config):
   match_threshold: float = 0.5
   unmatched_threshold: float = 0.5
   aug_rand_hflip: bool = False
-  aug_rand_jpeg: common.RandJpegQuality | None = None
+  aug_rand_jpeg: Optional[common.RandJpegQuality] = None
   aug_scale_min: float = 1.0
   aug_scale_max: float = 1.0
   skip_crowd_during_training: bool = True

official/vision/dataloaders/retinanet_input.py

@@ -18,7 +18,7 @@
 into (image, labels) tuple for RetinaNet.
 """
 
-from typing import Optional
+from typing import Optional, List
 
 from absl import logging
 import tensorflow as tf, tf_keras
@@ -37,17 +37,17 @@ class Parser(parser.Parser):
 
   def __init__(self,
                output_size,
-               min_level: int | None,
+               min_level: Optional[int] = None,
                max_level,
-               num_scales: int | None,
-               aspect_ratios: list[float] | None,
-               anchor_size: float | None,
+               num_scales: Optional[int] = None,
+               aspect_ratios: Optional[List[float]] = None,
+               anchor_size: Optional[float] = None,
                match_threshold=0.5,
                unmatched_threshold=0.5,
                box_coder_weights=None,
                aug_type=None,
                aug_rand_hflip=False,
-               aug_rand_jpeg: cfg.RandJpegQuality | None = None,
+               aug_rand_jpeg: Optional[cfg.RandJpegQuality] = None,
                aug_scale_min=1.0,
                aug_scale_max=1.0,
                use_autoaugment=False,

official/vision/modeling/backbones/mobilenet.py

@@ -15,7 +15,7 @@
 """Contains definitions of MobileNet Networks."""
 
 import dataclasses
-from typing import Any
+from typing import Any, Optional
 
 from absl import logging
 import tensorflow as tf, tf_keras
@@ -90,9 +90,9 @@ class BlockSpec(hyperparams.Config):
   use_normalization: bool = True
   activation: str = 'relu6'
   # Used for block type InvertedResConv.
-  expand_ratio: float | None = 6.0
+  expand_ratio: Optional[float] = 6.0
   # Used for block type InvertedResConv with SE.
-  se_ratio: float | None = None
+  se_ratio: Optional[float] = None
   use_depthwise: bool = True
   use_residual: bool = True
   is_output: bool = True
@@ -143,8 +143,8 @@ def __init__(
       use_explicit_padding: bool = False,
       activation: str = 'relu6',
       kernel_initializer: str = 'VarianceScaling',
-      kernel_regularizer: tf_keras.regularizers.Regularizer | None = None,
-      bias_regularizer: tf_keras.regularizers.Regularizer | None = None,
+      kernel_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
       use_normalization: bool = True,
       use_sync_bn: bool = False,
       norm_momentum: float = 0.99,
@@ -1228,10 +1228,10 @@ def __init__(
       norm_momentum: float = 0.99,
       norm_epsilon: float = 0.001,
       kernel_initializer: str = 'VarianceScaling',
-      kernel_regularizer: tf_keras.regularizers.Regularizer | None = None,
-      bias_regularizer: tf_keras.regularizers.Regularizer | None = None,
+      kernel_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
       # The followings should be kept the same most of the times.
-      output_stride: int | None = None,
+      output_stride: Optional[int] = None,
       min_depth: int = 8,
       # divisible is not used in MobileNetV1.
       divisible_by: int = 8,
@@ -1593,7 +1593,7 @@ def build_mobilenet(
     input_specs: tf_keras.layers.InputSpec,
     backbone_config: hyperparams.Config,
     norm_activation_config: hyperparams.Config,
-    l2_regularizer: tf_keras.regularizers.Regularizer | None = None,
+    l2_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
 ) -> tf_keras.Model:
   """Builds MobileNet backbone from a config."""
   backbone_type = backbone_config.type

official/vision/modeling/factory.py

@@ -14,7 +14,7 @@
 
 """Factory methods to build models."""
 
-from typing import Mapping, Optional
+from typing import Mapping, Optional, Union, Dict
 
 import tensorflow as tf, tf_keras
 
@@ -263,8 +263,8 @@ def build_retinanet(
     l2_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
     backbone: Optional[tf_keras.Model] = None,
     decoder: Optional[tf_keras.Model] = None,
-    num_anchors_per_location: int | dict[str, int] | None = None,
-    anchor_boxes: Mapping[str, tf.Tensor] | None = None,
+    num_anchors_per_location: Optional[Union[int, Dict[str, int]]] = None,
+    anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
 ) -> tf_keras.Model:
   """Builds a RetinaNet model.
 

official/vision/modeling/heads/dense_prediction_heads.py

@@ -31,7 +31,7 @@ def __init__(
       min_level: int,
       max_level: int,
       num_classes: int,
-      num_anchors_per_location: int | dict[str, int],
+      num_anchors_per_location: Union[int, Dict[str, int]] = None,
       num_convs: int = 4,
       num_filters: int = 256,
       attribute_heads: Optional[List[Dict[str, Any]]] = None,

official/vision/modeling/layers/detection_generator.py

@@ -792,7 +792,7 @@ def _generate_detections_tflite(
     raw_scores: Mapping[str, tf.Tensor],
     anchor_boxes: Mapping[str, tf.Tensor],
     config: Dict[str, Any],
-    box_coder_weights: List[float] | None = None,
+    box_coder_weights: Optional[List[float]] = None,
 ) -> Sequence[Any]:
   """Generate detections for conversion to TFLite.
 

official/vision/modeling/layers/nn_blocks.py

@@ -787,12 +787,12 @@ def __init__(
       start_dw_kernel_size: int = 0,
       middle_dw_kernel_size: int = 3,
       end_dw_kernel_size: int = 0,
-      stochastic_depth_drop_rate: float | None = None,
+      stochastic_depth_drop_rate: Optional[float] = None,
       kernel_initializer: str = 'VarianceScaling',
-      kernel_regularizer: tf_keras.regularizers.Regularizer | None = None,
-      bias_regularizer: tf_keras.regularizers.Regularizer | None = None,
+      kernel_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
       activation: str = 'relu',
-      depthwise_activation: str | None = None,
+      depthwise_activation: Optional[str] = None,
       use_sync_bn: bool = False,
       dilation_rate: int = 1,
       divisible_by: int = 1,

official/vision/modeling/retinanet_model.py

@@ -30,7 +30,7 @@ def __init__(self,
                decoder: tf_keras.Model,
                head: tf_keras.layers.Layer,
                detection_generator: tf_keras.layers.Layer,
-               anchor_boxes: Mapping[str, tf.Tensor] | None = None,
+               anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
                min_level: Optional[int] = None,
                max_level: Optional[int] = None,
                num_scales: Optional[int] = None,
@@ -81,14 +81,13 @@ def __init__(self,
     self._detection_generator = detection_generator
     self._anchor_boxes = anchor_boxes
 
-  def call(
-      self,  # pytype: disable=annotation-type-mismatch
-      images: Union[tf.Tensor, Sequence[tf.Tensor]],
-      image_shape: Optional[tf.Tensor] = None,
-      anchor_boxes: Mapping[str, tf.Tensor] | None = None,
-      output_intermediate_features: bool = False,
-      training: bool = None,
-  ) -> Mapping[str, tf.Tensor]:
+
+  def call(self,
+           images: Union[tf.Tensor, Sequence[tf.Tensor]],
+           image_shape: Optional[tf.Tensor] = None,
+           anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
+           output_intermediate_features: bool = False,
+           training: bool = None) -> Mapping[str, tf.Tensor]:
     """Forward pass of the RetinaNet model.
 
     Args:
@@ -241,7 +240,7 @@ def detection_generator(self) -> tf_keras.layers.Layer:
     return self._detection_generator
 
   @property
-  def anchor_boxes(self) -> Mapping[str, tf.Tensor] | None:
+  def anchor_boxes(self) -> Optional[Mapping[str, tf.Tensor]]:
     return self._anchor_boxes
 
   def get_config(self) -> Mapping[str, Any]:

official/vision/ops/augment.py

@@ -2856,7 +2856,7 @@ class SSDRandomCrop(ImageAugment):
 
   def __init__(
       self,
-      params: Sequence[configs.SSDRandomCropParam] | None = None,
+      params: Union[Sequence[configs.SSDRandomCropParam], None] = None,
       aspect_ratio_range: tuple[float, float] = (0.5, 2.0),
       area_range: tuple[float, float] = (0.1, 1.0),
   ):