Skip to content
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions official/nlp/modeling/layers/transformer_encoder_block.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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.

Expand All @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"""Defines base abstract uplift network layers."""

import abc
from typing import Union, Optional

import tensorflow as tf, tf_keras

Expand All @@ -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()
4 changes: 3 additions & 1 deletion official/recommendation/uplift/metrics/label_mean.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down
3 changes: 2 additions & 1 deletion official/recommendation/uplift/metrics/label_variance.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.

Expand Down
10 changes: 5 additions & 5 deletions official/recommendation/uplift/metrics/metric_configs.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,12 @@

from collections.abc import Mapping
import dataclasses
from typing import Any
from typing import Any, Optional

from official.core.config_definitions import base_config


@dataclasses.dataclass(kw_only=True)
@dataclasses.dataclass
class SlicedMetricConfig(base_config.Config):
"""Sliced metric configuration.

Expand All @@ -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]
Expand Down
9 changes: 5 additions & 4 deletions official/recommendation/uplift/metrics/sliced_metric.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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,
):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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.

Expand Down
4 changes: 3 additions & 1 deletion official/recommendation/uplift/metrics/uplift_mean.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down
6 changes: 3 additions & 3 deletions official/recommendation/uplift/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion official/vision/configs/retinanet.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 6 additions & 6 deletions official/vision/dataloaders/retinanet_input.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
into (image, labels) tuple for RetinaNet.
"""

from typing import Optional
from typing import Optional, List

# Import libraries

Expand All @@ -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: 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: cfg.RandJpegQuality | None = None,
aug_rand_jpeg: Optional[cfg.RandJpegQuality] = None,
aug_scale_min=1.0,
aug_scale_max=1.0,
use_autoaugment=False,
Expand Down
18 changes: 9 additions & 9 deletions official/vision/modeling/backbones/mobilenet.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
"""Contains definitions of MobileNet Networks."""

import dataclasses
from typing import Any
from typing import Any, Optional

# Import libraries

Expand Down Expand Up @@ -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: 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
Expand Down Expand Up @@ -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: 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,
Expand Down Expand Up @@ -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: 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,
Expand Down Expand Up @@ -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: Optional[tf_keras.regularizers.Regularizer] = None,
) -> tf_keras.Model:
"""Builds MobileNet backbone from a config."""
backbone_type = backbone_config.type
Expand Down
6 changes: 3 additions & 3 deletions official/vision/modeling/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.

Expand Down
2 changes: 1 addition & 1 deletion official/vision/modeling/heads/dense_prediction_heads.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion official/vision/modeling/layers/detection_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: Optional[List[float]] = None,
) -> Sequence[Any]:
"""Generate detections for conversion to TFLite.

Expand Down
8 changes: 4 additions & 4 deletions official/vision/modeling/layers/nn_blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: 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,
Expand Down
6 changes: 3 additions & 3 deletions official/vision/modeling/retinanet_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: Optional[Mapping[str, tf.Tensor]] = None,
min_level: Optional[int] = None,
max_level: Optional[int] = None,
num_scales: Optional[int] = None,
Expand Down Expand Up @@ -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: Optional[Mapping[str, tf.Tensor]] = None,
output_intermediate_features: bool = False,
training: bool = None) -> Mapping[str, tf.Tensor]:
"""Forward pass of the RetinaNet model.
Expand Down Expand Up @@ -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) -> Optional[Mapping[str, tf.Tensor]]:
return self._anchor_boxes

def get_config(self) -> Mapping[str, Any]:
Expand Down
2 changes: 1 addition & 1 deletion official/vision/ops/augment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
):
Expand Down
6 changes: 3 additions & 3 deletions official/vision/ops/preprocess_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
Loading
Loading