wehub-resource-sync
82 lines
3.0 KiB
Python
82 lines
3.0 KiB
Python
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""TPU hardware feature info."""
|
|
import enum
|
|
from tensorflow.core.protobuf.tpu import topology_pb2
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
|
|
|
|
@tf_export("tpu.experimental.HardwareFeature")
|
|
class HardwareFeature(object):
|
|
"""class holds all the feature info about the TPU."""
|
|
|
|
def __init__(self, tpu_hardware_feature_proto):
|
|
"""Store TPU hardware feature info.
|
|
|
|
Args:
|
|
tpu_hardware_feature_proto: protobuf which describe the tpu hardware
|
|
feature.
|
|
"""
|
|
self.tpu_hardware_feature_proto = tpu_hardware_feature_proto
|
|
|
|
class EmbeddingFeature(enum.Enum):
|
|
"""Embedding feature flag strings.
|
|
|
|
UNSUPPORTED: No embedding lookup accelerator available on the tpu.
|
|
V1: Embedding lookup accelerator V1. The embedding lookup operation can only
|
|
be placed at the beginning of computation. Only one instance of
|
|
embedding
|
|
lookup layer is allowed.
|
|
V2: Embedding lookup accelerator V2. The embedding lookup operation can be
|
|
placed anywhere of the computation. Multiple instances of embedding
|
|
lookup layer is allowed.
|
|
"""
|
|
UNSUPPORTED = "UNSUPPORTED"
|
|
V1 = "V1"
|
|
V2 = "V2"
|
|
|
|
@classmethod
|
|
def _embedding_feature_proto_to_string(cls, embedding_feature_proto):
|
|
"""Convert the embedding feature proto to enum string."""
|
|
embedding_feature_proto_to_string_map = {
|
|
topology_pb2.TPUHardwareFeature.EmbeddingFeature.UNSUPPORTED:
|
|
HardwareFeature.EmbeddingFeature.UNSUPPORTED,
|
|
topology_pb2.TPUHardwareFeature.EmbeddingFeature.V1:
|
|
HardwareFeature.EmbeddingFeature.V1,
|
|
topology_pb2.TPUHardwareFeature.EmbeddingFeature.V2:
|
|
HardwareFeature.EmbeddingFeature.V2
|
|
}
|
|
return embedding_feature_proto_to_string_map.get(
|
|
embedding_feature_proto, HardwareFeature.EmbeddingFeature.UNSUPPORTED)
|
|
|
|
@property
|
|
def embedding_feature(self):
|
|
"""TPU embedding feature.
|
|
|
|
Returns:
|
|
An EmbeddingFeature enum.
|
|
"""
|
|
return HardwareFeature._embedding_feature_proto_to_string(
|
|
self.tpu_hardware_feature_proto.embedding_feature)
|
|
|
|
@property
|
|
def num_embedding_devices_per_chip(self):
|
|
"""Number of embedding accelerator devices per chip.
|
|
|
|
Returns:
|
|
Number of embedding devices per chip.
|
|
"""
|
|
return self.tpu_hardware_feature_proto.num_embedding_devices_per_chip
|