From c6031a88d6989cb10d8a8c92e54dcf844d5e3e52 Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Fri, 13 Sep 2024 06:12:30 +0000
Subject: [PATCH 1/7] Make some api's compatible with python3.7+

---
 .../layers/transformer_encoder_block.py        |  6 +++---
 .../uplift_networks/base_uplift_networks.py    |  6 ++++--
 .../uplift/metrics/label_mean.py               |  3 ++-
 .../uplift/metrics/label_variance.py           |  3 ++-
 .../uplift/metrics/metric_configs.py           | 10 +++++-----
 .../uplift/metrics/sliced_metric.py            | 10 ++++++----
 .../uplift/metrics/treatment_fraction.py       |  3 ++-
 .../uplift/metrics/uplift_mean.py              |  3 ++-
 official/recommendation/uplift/types.py        |  5 +++--
 official/vision/configs/retinanet.py           |  2 +-
 official/vision/dataloaders/retinanet_input.py | 12 ++++++------
 .../vision/modeling/backbones/mobilenet.py     | 18 +++++++++---------
 official/vision/modeling/factory.py            |  6 +++---
 .../modeling/heads/dense_prediction_heads.py   |  2 +-
 .../modeling/layers/detection_generator.py     |  4 ++--
 official/vision/modeling/layers/nn_blocks.py   |  8 ++++----
 official/vision/modeling/retinanet_model.py    |  6 +++---
 official/vision/ops/augment.py                 |  2 +-
 official/vision/ops/preprocess_ops.py          |  6 +++---
 .../vision/serving/export_saved_model_lib.py   |  2 +-
 20 files changed, 63 insertions(+), 54 deletions(-)

diff --git a/official/nlp/modeling/layers/transformer_encoder_block.py b/official/nlp/modeling/layers/transformer_encoder_block.py
index 20f2280543d..a50a0911ecf 100644
--- 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
index 0cce8a459ba..64552ab9828 100644
--- 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
index a86cd7f29cb..31b48bde31d 100644
--- 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
index e198bcdedf4..fd6e78ca620 100644
--- 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
index 05f3898132e..e07cda4018f 100644
--- 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
index cae1cfa341a..ef5152e2a41 100644
--- 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
index 6fc5d54735a..fb1d9a62c49 100644
--- 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
index 946adbc0f05..31d1e860ebb 100644
--- 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
index 8e8ce450c94..3d21047ee1a 100644
--- 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
index 5c04e5693d8..feff3c39423 100644
--- 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
index e21f2c6ebfd..c720f258857 100644
--- 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
index 0929b1775b0..75c8abae11e 100644
--- 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
index 56e67b21a93..cbca87a66b2 100644
--- 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
index 9a905e69874..eb4a8824de6 100644
--- 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
index 6cab54d2a15..6929aa45490 100644
--- 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
index 02c214a40f2..c72b85600ca 100644
--- 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
index 7127b1e6f92..32d57f4fa3c 100644
--- 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
index 90c7266d4de..680bdc9a34c 100644
--- 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
index ddf23c56386..4e296a1abdd 100644
--- 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
index c6f1b057ea4..f094fd4d042 100644
--- a/official/vision/serving/export_saved_model_lib.py
+++ b/official/vision/serving/export_saved_model_lib.py
@@ -35,8 +35,8 @@ def export_inference_graph(
     batch_size: Optional[int],
     input_image_size: List[int],
     params: cfg.ExperimentConfig,
-    checkpoint_path: str | None,
     export_dir: str,
+    checkpoint_path: Union[str, None] = None,
     num_channels: Optional[int] = 3,
     export_module: Optional[export_base.ExportModule] = None,
     export_checkpoint_subdir: Optional[str] = None,

From 5401505ac4ed921b7d36b64338bb9e7556cba7ff Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Mon, 16 Sep 2024 08:29:32 +0000
Subject: [PATCH 2/7] line33 syntax changes

---
 official/nlp/modeling/layers/transformer_encoder_block.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/official/nlp/modeling/layers/transformer_encoder_block.py b/official/nlp/modeling/layers/transformer_encoder_block.py
index a50a0911ecf..86f497b4862 100644
--- a/official/nlp/modeling/layers/transformer_encoder_block.py
+++ b/official/nlp/modeling/layers/transformer_encoder_block.py
@@ -30,7 +30,7 @@ def __init__(
       self,
       axis: Union[int , Sequence[int]] = -1,
       epsilon: float = 1e-6,
-      **kwargs,
+      **kwargs
   ):
     """Initializes RMSNorm.
 

From 441f63f7779d3e97c267b8b53fa33e025f2f2f37 Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Tue, 8 Oct 2024 05:37:34 +0000
Subject: [PATCH 3/7] lint changes in few files

---
 official/nlp/modeling/layers/transformer_encoder_block.py    | 5 +++--
 .../uplift/layers/uplift_networks/base_uplift_networks.py    | 3 +--
 official/recommendation/uplift/metrics/label_mean.py         | 3 ++-
 official/recommendation/uplift/metrics/label_variance.py     | 2 +-
 official/recommendation/uplift/metrics/metric_configs.py     | 2 +-
 official/recommendation/uplift/metrics/sliced_metric.py      | 3 +--
 official/recommendation/uplift/metrics/uplift_mean.py        | 3 ++-
 official/recommendation/uplift/types.py                      | 3 +--
 8 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/official/nlp/modeling/layers/transformer_encoder_block.py b/official/nlp/modeling/layers/transformer_encoder_block.py
index 86f497b4862..a8abc953df7 100644
--- a/official/nlp/modeling/layers/transformer_encoder_block.py
+++ b/official/nlp/modeling/layers/transformer_encoder_block.py
@@ -28,7 +28,7 @@ class RMSNorm(tf_keras.layers.Layer):
 
   def __init__(
       self,
-      axis: Union[int , Sequence[int]] = -1,
+      axis: Union[int, Sequence[int]] = -1,
       epsilon: float = 1e-6,
       **kwargs
   ):
@@ -43,7 +43,8 @@ def __init__(
     self.axis = [axis] if isinstance(axis, int) else axis
     self.epsilon = epsilon
 
-  def build(self, input_shape: Union[tf.TensorShape, Sequence[Union[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
index 64552ab9828..907ab663f42 100644
--- a/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py
+++ b/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py
@@ -15,13 +15,12 @@
 """Defines base abstract uplift network layers."""
 
 import abc
+from typing import Union
 
 import tensorflow as tf, tf_keras
 
 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.
diff --git a/official/recommendation/uplift/metrics/label_mean.py b/official/recommendation/uplift/metrics/label_mean.py
index 31b48bde31d..aacaa0ef96a 100644
--- a/official/recommendation/uplift/metrics/label_mean.py
+++ b/official/recommendation/uplift/metrics/label_mean.py
@@ -14,12 +14,13 @@
 
 """Keras metric for computing the label mean sliced by treatment group."""
 
+from typing import Union
+
 import tensorflow as tf, tf_keras
 
 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):
diff --git a/official/recommendation/uplift/metrics/label_variance.py b/official/recommendation/uplift/metrics/label_variance.py
index fd6e78ca620..e9db2f45c03 100644
--- a/official/recommendation/uplift/metrics/label_variance.py
+++ b/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 Union
 
 import tensorflow as tf, tf_keras
 
@@ -20,7 +21,6 @@
 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):
diff --git a/official/recommendation/uplift/metrics/metric_configs.py b/official/recommendation/uplift/metrics/metric_configs.py
index e07cda4018f..0487ccc04d4 100644
--- a/official/recommendation/uplift/metrics/metric_configs.py
+++ b/official/recommendation/uplift/metrics/metric_configs.py
@@ -35,7 +35,7 @@ class SlicedMetricConfig(base_config.Config):
 
   slicing_feature: Union[str, None] = None
   slicing_spec: Union[Mapping[str, int], None] = None
-  slicing_feature_dtype: Union[str, 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
index ef5152e2a41..ac0bf7dd5e0 100644
--- a/official/recommendation/uplift/metrics/sliced_metric.py
+++ b/official/recommendation/uplift/metrics/sliced_metric.py
@@ -15,11 +15,10 @@
 """Keras metric for reporting metrics sliced by a feature."""
 
 import copy
+from typing import Union
 
 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.
diff --git a/official/recommendation/uplift/metrics/uplift_mean.py b/official/recommendation/uplift/metrics/uplift_mean.py
index 31d1e860ebb..ee94f91d682 100644
--- a/official/recommendation/uplift/metrics/uplift_mean.py
+++ b/official/recommendation/uplift/metrics/uplift_mean.py
@@ -14,12 +14,13 @@
 
 """Keras metric for computing the mean uplift sliced by treatment group."""
 
+from typing import Union
+
 import tensorflow as tf, tf_keras
 
 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):
diff --git a/official/recommendation/uplift/types.py b/official/recommendation/uplift/types.py
index 3d21047ee1a..95f316fc6e5 100644
--- a/official/recommendation/uplift/types.py
+++ b/official/recommendation/uplift/types.py
@@ -13,12 +13,11 @@
 # limitations under the License.
 
 """Defines types used by the keras uplift modeling library."""
+from typing import Union
 
 import tensorflow as tf, tf_keras
-from typing import Union
 
 TensorType = Union[tf.Tensor, tf.SparseTensor, tf.RaggedTensor]
-
 ListOfTensors = list[TensorType]
 TupleOfTensors = tuple[TensorType, ...]
 DictOfTensors = dict[str, TensorType]

From 001ec92b3574896d6e9f65b8a5b46afc62e5f9e6 Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Mon, 28 Oct 2024 04:42:53 +0000
Subject: [PATCH 4/7] Changed redundant Union(X, None) to Optional(X) in
 related files.

---
 .../layers/transformer_encoder_block.py        |  2 +-
 .../uplift_networks/base_uplift_networks.py    |  6 +++---
 .../uplift/metrics/label_mean.py               |  4 ++--
 .../uplift/metrics/label_variance.py           |  4 ++--
 .../uplift/metrics/metric_configs.py           |  8 ++++----
 .../uplift/metrics/sliced_metric.py            |  8 ++++----
 .../uplift/metrics/treatment_fraction.py       |  4 ++--
 .../uplift/metrics/uplift_mean.py              |  4 ++--
 official/vision/configs/retinanet.py           |  2 +-
 official/vision/dataloaders/retinanet_input.py | 12 ++++++------
 .../vision/modeling/backbones/mobilenet.py     | 18 +++++++++---------
 official/vision/modeling/factory.py            |  2 +-
 .../modeling/layers/detection_generator.py     |  4 ++--
 official/vision/modeling/layers/nn_blocks.py   |  8 ++++----
 official/vision/modeling/retinanet_model.py    |  6 +++---
 .../vision/serving/export_saved_model_lib.py   |  2 +-
 16 files changed, 47 insertions(+), 47 deletions(-)

diff --git a/official/nlp/modeling/layers/transformer_encoder_block.py b/official/nlp/modeling/layers/transformer_encoder_block.py
index a8abc953df7..1c5de85a892 100644
--- a/official/nlp/modeling/layers/transformer_encoder_block.py
+++ b/official/nlp/modeling/layers/transformer_encoder_block.py
@@ -44,7 +44,7 @@ def __init__(
     self.epsilon = epsilon
 
   def build(self, 
-            input_shape: Union[tf.TensorShape, Sequence[Union[int, None]]]):
+            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:
diff --git a/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py b/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py
index 907ab663f42..d74ecc2cad3 100644
--- a/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py
+++ b/official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py
@@ -15,7 +15,7 @@
 """Defines base abstract uplift network layers."""
 
 import abc
-from typing import Union
+from typing import Union, Optional
 
 import tensorflow as tf, tf_keras
 
@@ -34,7 +34,7 @@ class BaseTwoTowerUpliftNetwork(tf_keras.layers.Layer, metaclass=abc.ABCMeta):
   def call(
       self,
       inputs: types.DictOfTensors,
-      training: Union[bool, None] = None,
-      mask: Union[tf.Tensor, None] = None,
+      training: Optional[bool] = None,
+      mask: Optional[tf.Tensor] = None,
   ) -> types.TwoTowerTrainingOutputs:
     raise NotImplementedError()
diff --git a/official/recommendation/uplift/metrics/label_mean.py b/official/recommendation/uplift/metrics/label_mean.py
index aacaa0ef96a..917aec59e1e 100644
--- a/official/recommendation/uplift/metrics/label_mean.py
+++ b/official/recommendation/uplift/metrics/label_mean.py
@@ -14,7 +14,7 @@
 
 """Keras metric for computing the label mean sliced by treatment group."""
 
-from typing import Union
+from typing import Union, Optional
 
 import tensorflow as tf, tf_keras
 
@@ -73,7 +73,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: Union[tf.Tensor, None] = None,
+      sample_weight: Optional[tf.Tensor] = 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
index e9db2f45c03..dcf6b417dc5 100644
--- a/official/recommendation/uplift/metrics/label_variance.py
+++ b/official/recommendation/uplift/metrics/label_variance.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Keras metric for computing the label variance sliced by treatment group."""
-from typing import Union
+from typing import Optional
 
 import tensorflow as tf, tf_keras
 
@@ -73,7 +73,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: Union[tf.Tensor, None] = None,
+      sample_weight: Optional[tf.Tensor] = 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
index 0487ccc04d4..4b5ecf3a8e8 100644
--- a/official/recommendation/uplift/metrics/metric_configs.py
+++ b/official/recommendation/uplift/metrics/metric_configs.py
@@ -16,7 +16,7 @@
 
 from collections.abc import Mapping
 import dataclasses
-from typing import Any, Union
+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: Union[str, None] = None
-  slicing_spec: Union[Mapping[str, int], None] = None
-  slicing_feature_dtype: Union[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]
diff --git a/official/recommendation/uplift/metrics/sliced_metric.py b/official/recommendation/uplift/metrics/sliced_metric.py
index ac0bf7dd5e0..e550aa49957 100644
--- a/official/recommendation/uplift/metrics/sliced_metric.py
+++ b/official/recommendation/uplift/metrics/sliced_metric.py
@@ -15,7 +15,7 @@
 """Keras metric for reporting metrics sliced by a feature."""
 
 import copy
-from typing import Union
+from typing import Union, Optional
 
 import tensorflow as tf, tf_keras
 
@@ -68,8 +68,8 @@ def __init__(
       self,
       metric: tf_keras.metrics.Metric,
       slicing_spec: Union[dict[str, str], dict[str, int]],
-      slicing_feature_dtype: Union[tf.DType, None] = None,
-      name: Union[str, None] = None,
+      slicing_feature_dtype: Optional[tf.DType] = None,
+      name: Optional[str] = None,
   ):
     """Initializes the instance.
 
@@ -124,7 +124,7 @@ def __init__(
   def update_state(
       self,
       *args: tf.Tensor,
-      sample_weight: Union[tf.Tensor, None] = None,
+      sample_weight: Optional[tf.Tensor] = None,
       slicing_feature: tf.Tensor,
       **kwargs,
   ):
diff --git a/official/recommendation/uplift/metrics/treatment_fraction.py b/official/recommendation/uplift/metrics/treatment_fraction.py
index fb1d9a62c49..c63bed8e6b5 100644
--- a/official/recommendation/uplift/metrics/treatment_fraction.py
+++ b/official/recommendation/uplift/metrics/treatment_fraction.py
@@ -18,7 +18,7 @@
 
 from official.recommendation.uplift import types
 
-from typing import Union
+from typing import Optional
 
 @tf_keras.utils.register_keras_serializable(package="Uplift")
 class TreatmentFraction(tf_keras.metrics.Metric):
@@ -58,7 +58,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: Union[tf.Tensor, None] = None,
+      sample_weight: Optional[tf.Tensor] = None,
   ) -> None:
     """Updates the treatment fraction.
 
diff --git a/official/recommendation/uplift/metrics/uplift_mean.py b/official/recommendation/uplift/metrics/uplift_mean.py
index ee94f91d682..e34b723aa22 100644
--- a/official/recommendation/uplift/metrics/uplift_mean.py
+++ b/official/recommendation/uplift/metrics/uplift_mean.py
@@ -14,7 +14,7 @@
 
 """Keras metric for computing the mean uplift sliced by treatment group."""
 
-from typing import Union
+from typing import Optional
 
 import tensorflow as tf, tf_keras
 
@@ -70,7 +70,7 @@ def update_state(
       self,
       y_true: tf.Tensor,
       y_pred: types.TwoTowerTrainingOutputs,
-      sample_weight: Union[tf.Tensor, None] = None,
+      sample_weight: Optional[tf.Tensor] = None,
   ) -> None:
     """Updates the overall, control and treatment uplift means.
 
diff --git a/official/vision/configs/retinanet.py b/official/vision/configs/retinanet.py
index feff3c39423..dd5e0ba15d0 100644
--- 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: Union[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
diff --git a/official/vision/dataloaders/retinanet_input.py b/official/vision/dataloaders/retinanet_input.py
index c720f258857..a3bf8f38320 100644
--- 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, Union, List
+from typing import Optional, List
 
 # Import libraries
 
@@ -40,16 +40,16 @@ class Parser(parser.Parser):
   def __init__(self,
                output_size,
                max_level,
-               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,
+               min_level: Optional[int] = 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: Union[cfg.RandJpegQuality, None] = None,
+               aug_rand_jpeg: Optional[cfg.RandJpegQuality] = 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
index 75c8abae11e..9521d587103 100644
--- 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, Union
+from typing import Any, Optional
 
 # Import libraries
 
@@ -92,9 +92,9 @@ class BlockSpec(hyperparams.Config):
   use_normalization: bool = True
   activation: str = 'relu6'
   # Used for block type InvertedResConv.
-  expand_ratio: Union[float, None] = 6.0
+  expand_ratio: Optional[float] = 6.0
   # Used for block type InvertedResConv with SE.
-  se_ratio: Union[float, None] = None
+  se_ratio: Optional[float] = 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: Union[tf_keras.regularizers.Regularizer, None] = None,
-      bias_regularizer: Union[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,
@@ -1230,10 +1230,10 @@ def __init__(
       norm_momentum: float = 0.99,
       norm_epsilon: float = 0.001,
       kernel_initializer: str = 'VarianceScaling',
-      kernel_regularizer: Union[tf_keras.regularizers.Regularizer, None] = None,
-      bias_regularizer: Union[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: Union[int, None] = None,
+      output_stride: Optional[int] = 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: Union[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
diff --git a/official/vision/modeling/factory.py b/official/vision/modeling/factory.py
index cbca87a66b2..2814bce015b 100644
--- a/official/vision/modeling/factory.py
+++ b/official/vision/modeling/factory.py
@@ -264,7 +264,7 @@ def build_retinanet(
     backbone: Optional[tf_keras.Model] = None,
     decoder: Optional[tf_keras.Model] = None,
     num_anchors_per_location: Union[int, Dict[str, int], None] = None,
-    anchor_boxes: Union[Mapping[str, tf.Tensor], None] = None,
+    anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
 ) -> tf_keras.Model:
   """Builds a RetinaNet model.
 
diff --git a/official/vision/modeling/layers/detection_generator.py b/official/vision/modeling/layers/detection_generator.py
index 6929aa45490..c9d248973f4 100644
--- 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, Union
+from typing import Any, Dict, List, Optional, Mapping, Sequence, Tuple
 
 # 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: Union[List[float], None] = None,
+    box_coder_weights: Optional[List[float]] = 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
index c72b85600ca..dea05d66c1f 100644
--- 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: Union[float, None] = None,
+      stochastic_depth_drop_rate: Optional[float] = None,
       kernel_initializer: str = 'VarianceScaling',
-      kernel_regularizer: Union[tf_keras.regularizers.Regularizer, None] = None,
-      bias_regularizer: Union[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: Union[str, None] = None,
+      depthwise_activation: Optional[str] = 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
index 32d57f4fa3c..cf721dcb83e 100644
--- 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: Union[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,
@@ -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: Union[Mapping[str, tf.Tensor], None] = 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.
@@ -240,7 +240,7 @@ def detection_generator(self) -> tf_keras.layers.Layer:
     return self._detection_generator
 
   @property
-  def anchor_boxes(self) -> Union[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]:
diff --git a/official/vision/serving/export_saved_model_lib.py b/official/vision/serving/export_saved_model_lib.py
index f094fd4d042..879b638d3df 100644
--- a/official/vision/serving/export_saved_model_lib.py
+++ b/official/vision/serving/export_saved_model_lib.py
@@ -36,7 +36,7 @@ def export_inference_graph(
     input_image_size: List[int],
     params: cfg.ExperimentConfig,
     export_dir: str,
-    checkpoint_path: Union[str, None] = None,
+    checkpoint_path: Optional[str] = None,
     num_channels: Optional[int] = 3,
     export_module: Optional[export_base.ExportModule] = None,
     export_checkpoint_subdir: Optional[str] = None,

From 276a37093f5bd38121d5d93b63f3fc7333e3f4ab Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Mon, 18 Nov 2024 08:28:55 +0000
Subject: [PATCH 5/7] Done the suggested changes in factory.py

---
 official/vision/modeling/factory.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/official/vision/modeling/factory.py b/official/vision/modeling/factory.py
index 2814bce015b..998894464c4 100644
--- a/official/vision/modeling/factory.py
+++ b/official/vision/modeling/factory.py
@@ -263,7 +263,7 @@ 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: Union[int, Dict[str, int], 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.

From 0b86c7d24a8c9ebd9e3a3ddf177203cd077d46c8 Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Thu, 21 Nov 2024 08:33:32 +0000
Subject: [PATCH 6/7] Done the suggested changes for keeping the original order
 of arguments

---
 official/vision/dataloaders/retinanet_input.py    | 2 +-
 official/vision/serving/export_saved_model_lib.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/official/vision/dataloaders/retinanet_input.py b/official/vision/dataloaders/retinanet_input.py
index a3bf8f38320..9875bab81f2 100644
--- a/official/vision/dataloaders/retinanet_input.py
+++ b/official/vision/dataloaders/retinanet_input.py
@@ -39,8 +39,8 @@ class Parser(parser.Parser):
 
   def __init__(self,
                output_size,
-               max_level,
                min_level: Optional[int] = None,
+               max_level,
                num_scales: Optional[int] = None,
                aspect_ratios: Optional[List[float]] = None,
                anchor_size: Optional[float] = None,
diff --git a/official/vision/serving/export_saved_model_lib.py b/official/vision/serving/export_saved_model_lib.py
index 879b638d3df..81f0161ec26 100644
--- a/official/vision/serving/export_saved_model_lib.py
+++ b/official/vision/serving/export_saved_model_lib.py
@@ -35,8 +35,8 @@ def export_inference_graph(
     batch_size: Optional[int],
     input_image_size: List[int],
     params: cfg.ExperimentConfig,
-    export_dir: str,
     checkpoint_path: Optional[str] = None,
+    export_dir: str,
     num_channels: Optional[int] = 3,
     export_module: Optional[export_base.ExportModule] = None,
     export_checkpoint_subdir: Optional[str] = None,

From 1e026a126712d4cda4fb2579e6c2e3ce313d463a Mon Sep 17 00:00:00 2001
From: bharatjetti <jettibharat@google.com>
Date: Wed, 18 Dec 2024 08:17:08 +0000
Subject: [PATCH 7/7] restored the 'kw_only=True'

---
 official/recommendation/uplift/metrics/metric_configs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/official/recommendation/uplift/metrics/metric_configs.py b/official/recommendation/uplift/metrics/metric_configs.py
index 4b5ecf3a8e8..1f704fc79a0 100644
--- a/official/recommendation/uplift/metrics/metric_configs.py
+++ b/official/recommendation/uplift/metrics/metric_configs.py
@@ -21,7 +21,7 @@
 from official.core.config_definitions import base_config
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(kw_only=True)
 class SlicedMetricConfig(base_config.Config):
   """Sliced metric configuration.