Make some api's compatible with python3.7+

bharatjetti · bharatjetti · commit c6031a88d698 · 2024-09-13T06:12:30.000Z
diff --git a/official/nlp/modeling/layers/transformer_encoder_block.py b/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,7 +28,7 @@ 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,
   ):
@@ -43,7 +43,7 @@ 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[Union[int, None]]]):
     input_shape = tf.TensorShape(input_shape)
     scale_shape = [1] * input_shape.rank
     for dim in self.axis:
diff --git a/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py b/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py
@@ -20,6 +20,8 @@
 
 from official.recommendation.uplift import types
 
+from typing import Union
+
 
 class BaseTwoTowerUpliftNetwork(tf_keras.layers.Layer, metaclass=abc.ABCMeta):
   """Abstract class for uplift layers that compute control and treatment logits.
@@ -33,7 +35,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: Union[bool, None] = None,
+      mask: Union[tf.Tensor, None] = None,
   ) -> types.TwoTowerTrainingOutputs:
     raise NotImplementedError()
diff --git a/official/recommendation/uplift/metrics/label_mean.py b/official/recommendation/uplift/metrics/label_mean.py
@@ -19,6 +19,7 @@
 from official.recommendation.uplift import types
 from official.recommendation.uplift.metrics import treatment_sliced_metric
 
+from typing import Union
 
 @tf_keras.utils.register_keras_serializable(package="Uplift")
 class LabelMean(tf_keras.metrics.Metric):
@@ -71,7 +72,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Union[tf.Tensor, None] = None,
   ):
     """Updates the overall, control and treatment label means.
 
diff --git a/official/recommendation/uplift/metrics/label_variance.py b/official/recommendation/uplift/metrics/label_variance.py
@@ -20,6 +20,7 @@
 from official.recommendation.uplift.metrics import treatment_sliced_metric
 from official.recommendation.uplift.metrics import variance
 
+from typing import Union
 
 @tf_keras.utils.register_keras_serializable(package="Uplift")
 class LabelVariance(tf_keras.metrics.Metric):
@@ -72,7 +73,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Union[tf.Tensor, None] = None,
   ):
     """Updates the overall, control and treatment label variances.
 
diff --git a/official/recommendation/uplift/metrics/metric_configs.py b/official/recommendation/uplift/metrics/metric_configs.py
@@ -16,12 +16,12 @@
 
 from collections.abc import Mapping
 import dataclasses
-from typing import Any
+from typing import Any, Union
 
 from official.core.config_definitions import base_config
 
 
-@dataclasses.dataclass(kw_only=True)
+@dataclasses.dataclass
 class SlicedMetricConfig(base_config.Config):
   """Sliced metric configuration.
 
@@ -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: Union[str, None] = None
+  slicing_spec: Union[Mapping[str, int], None] = None
+  slicing_feature_dtype: Union[str, None ]= None
 
   def __post_init__(
       self, default_params: dict[str, Any], restrictions: list[str]
diff --git a/official/recommendation/uplift/metrics/sliced_metric.py b/official/recommendation/uplift/metrics/sliced_metric.py
@@ -18,6 +18,8 @@
 
 import tensorflow as tf, tf_keras
 
+from typing import Union
+
 
 class SlicedMetric(tf_keras.metrics.Metric):
   """A metric sliced by integer, boolean, or string features.
@@ -66,9 +68,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: Union[tf.DType, None] = None,
+      name: Union[str, None] = None,
   ):
     """Initializes the instance.
 
@@ -123,7 +125,7 @@ def __init__(
   def update_state(
       self,
       *args: tf.Tensor,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Union[tf.Tensor, None] = None,
       slicing_feature: tf.Tensor,
       **kwargs,
   ):
diff --git a/official/recommendation/uplift/metrics/treatment_fraction.py b/official/recommendation/uplift/metrics/treatment_fraction.py
@@ -18,6 +18,7 @@
 
 from official.recommendation.uplift import types
 
+from typing import Union
 
 @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: Union[tf.Tensor, None] = None,
   ) -> None:
     """Updates the treatment fraction.
 
diff --git a/official/recommendation/uplift/metrics/uplift_mean.py b/official/recommendation/uplift/metrics/uplift_mean.py
@@ -19,6 +19,7 @@
 from official.recommendation.uplift import types
 from official.recommendation.uplift.metrics import treatment_sliced_metric
 
+from typing import Union
 
 @tf_keras.utils.register_keras_serializable(package="Uplift")
 class UpliftMean(tf_keras.metrics.Metric):
@@ -68,7 +69,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: tf.Tensor | None = None,
+      sample_weight: Union[tf.Tensor, None] = None,
   ) -> None:
     """Updates the overall, control and treatment uplift means.
 
diff --git a/official/recommendation/uplift/types.py b/official/recommendation/uplift/types.py
@@ -15,14 +15,15 @@
 """Defines types used by the keras uplift modeling library."""
 
 import tensorflow as tf, tf_keras
+from typing import Union
 
-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):
diff --git a/official/vision/configs/retinanet.py b/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: Union[common.RandJpegQuality, None] = None
   aug_scale_min: float = 1.0
   aug_scale_max: float = 1.0
   skip_crowd_during_training: bool = True
diff --git a/official/vision/dataloaders/retinanet_input.py b/official/vision/dataloaders/retinanet_input.py
@@ -18,7 +18,7 @@
 into (image, labels) tuple for RetinaNet.
 """
 
-from typing import Optional
+from typing import Optional, Union, List
 
 # Import libraries
 
@@ -39,17 +39,17 @@ class Parser(parser.Parser):
 
   def __init__(self,
                output_size,
-               min_level: int | None,
                max_level,
-               num_scales: int | None,
-               aspect_ratios: list[float] | None,
-               anchor_size: float | None,
+               min_level: Union[int, None] = None,
+               num_scales: Union[int, None] = None,
+               aspect_ratios: Union[List[float], None] = None,
+               anchor_size: Union[float, None] = 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: Union[cfg.RandJpegQuality, None] = None,
                aug_scale_min=1.0,
                aug_scale_max=1.0,
                use_autoaugment=False,
diff --git a/official/vision/modeling/backbones/mobilenet.py b/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, Union
 
 # Import libraries
 
@@ -92,9 +92,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: Union[float, None] = 6.0
   # Used for block type InvertedResConv with SE.
-  se_ratio: float | None = None
+  se_ratio: Union[float, None] = None
   use_depthwise: bool = True
   use_residual: bool = True
   is_output: bool = True
@@ -145,8 +145,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: Union[tf_keras.regularizers.Regularizer, None] = None,
+      bias_regularizer: Union[tf_keras.regularizers.Regularizer, None] = None,
       use_normalization: bool = True,
       use_sync_bn: bool = False,
       norm_momentum: float = 0.99,
@@ -1230,10 +1230,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: Union[tf_keras.regularizers.Regularizer, None] = None,
+      bias_regularizer: Union[tf_keras.regularizers.Regularizer, None] = None,
       # The followings should be kept the same most of the times.
-      output_stride: int | None = None,
+      output_stride: Union[int, None] = None,
       min_depth: int = 8,
       # divisible is not used in MobileNetV1.
       divisible_by: int = 8,
@@ -1595,7 +1595,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: Union[tf_keras.regularizers.Regularizer, None] = None,
 ) -> tf_keras.Model:
   """Builds MobileNet backbone from a config."""
   backbone_type = backbone_config.type
diff --git a/official/vision/modeling/factory.py b/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: Union[int, Dict[str, int], None] = None,
+    anchor_boxes: Union[Mapping[str, tf.Tensor], None] = None,
 ) -> tf_keras.Model:
   """Builds a RetinaNet model.
 
diff --git a/official/vision/modeling/heads/dense_prediction_heads.py b/official/vision/modeling/heads/dense_prediction_heads.py
@@ -33,7 +33,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,
diff --git a/official/vision/modeling/layers/detection_generator.py b/official/vision/modeling/layers/detection_generator.py
@@ -14,7 +14,7 @@
 
 """Contains definitions of generators to generate the final detections."""
 import contextlib
-from typing import Any, Dict, List, Optional, Mapping, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Mapping, Sequence, Tuple, Union
 
 # Import libraries
 
@@ -794,7 +794,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: Union[List[float], None] = None,
 ) -> Sequence[Any]:
   """Generate detections for conversion to TFLite.
 
diff --git a/official/vision/modeling/layers/nn_blocks.py b/official/vision/modeling/layers/nn_blocks.py
@@ -788,12 +788,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: Union[float, None] = None,
       kernel_initializer: str = 'VarianceScaling',
-      kernel_regularizer: tf_keras.regularizers.Regularizer | None = None,
-      bias_regularizer: tf_keras.regularizers.Regularizer | None = None,
+      kernel_regularizer: Union[tf_keras.regularizers.Regularizer, None] = None,
+      bias_regularizer: Union[tf_keras.regularizers.Regularizer, None] = None,
       activation: str = 'relu',
-      depthwise_activation: str | None = None,
+      depthwise_activation: Union[str, None] = None,
       use_sync_bn: bool = False,
       dilation_rate: int = 1,
       divisible_by: int = 1,
diff --git a/official/vision/modeling/retinanet_model.py b/official/vision/modeling/retinanet_model.py
@@ -31,7 +31,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: Union[Mapping[str, tf.Tensor], None] = None,
                min_level: Optional[int] = None,
                max_level: Optional[int] = None,
                num_scales: Optional[int] = None,
@@ -85,7 +85,7 @@ def __init__(self,
   def call(self,
            images: Union[tf.Tensor, Sequence[tf.Tensor]],
            image_shape: Optional[tf.Tensor] = None,
-           anchor_boxes: Mapping[str, tf.Tensor] | None = None,
+           anchor_boxes: Union[Mapping[str, tf.Tensor], None] = None,
            output_intermediate_features: bool = False,
            training: bool = None) -> Mapping[str, tf.Tensor]:
     """Forward pass of the RetinaNet model.
@@ -240,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) -> Union[Mapping[str, tf.Tensor], None]:
     return self._anchor_boxes
 
   def get_config(self) -> Mapping[str, Any]:
diff --git a/official/vision/ops/augment.py b/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),
   ):
diff --git a/official/vision/ops/preprocess_ops.py b/official/vision/ops/preprocess_ops.py
@@ -1217,9 +1217,9 @@ def random_crop(
 
 def random_jpeg_quality(
     image: tf.Tensor,
-    min_quality: int | tf.Tensor = 20,
-    max_quality: int | tf.Tensor = 100,
-    prob_to_apply: float | tf.Tensor = 0.6,
+    min_quality: Union[int, tf.Tensor] = 20,
+    max_quality: Union[int, tf.Tensor] = 100,
+    prob_to_apply: Union[float, tf.Tensor] = 0.6,
 ) -> tf.Tensor:
   """Randomly encode the image as jpeg and decode it.
 
diff --git a/official/vision/serving/export_saved_model_lib.py b/official/vision/serving/export_saved_model_lib.py
