def test_tensorflow2_datatypes():
# _NP_TO_TF contains all the mappings
# of numpy to tf types
try:
from tensorflow.python import _pywrap_bfloat16
# TF 2.x.x Implements a Custom Numpy Datatype for Brain Floating Type
# Which is currently only supported on TPUs
_np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
_NP_TO_TF.pop(_np_bfloat16)
except (ModuleNotFoundError, ValueError, ImportError):
pass
for _type in _NP_TO_TF:
try:
_get_proto_dtype(np.dtype(_type))
except Exception:
assert False, f"{_type} not supported"
assert True
from __future__ import division
from __future__ import print_function
import numpy as np
from six.moves import builtins
from tensorflow.core.framework import types_pb2
# We need to import pywrap_tensorflow prior to the bfloat wrapper to avoid
# protobuf errors where a file is defined twice on MacOS.
# pylint: disable=invalid-import-order,g-bad-import-order
from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import
from tensorflow.python import _pywrap_bfloat16
from tensorflow.python import _dtypes
from tensorflow.python.util.tf_export import tf_export
_np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
# pylint: disable=slots-on-old-class
@tf_export("dtypes.DType", "DType")
class DType(_dtypes.DType):
"""Represents the type of the elements in a `Tensor`.
The following `DType` objects are defined:
* `tf.float16`: 16-bit half-precision floating-point.
* `tf.float32`: 32-bit single-precision floating-point.
* `tf.float64`: 64-bit double-precision floating-point.
* `tf.bfloat16`: 16-bit truncated floating-point.
* `tf.complex64`: 64-bit single-precision complex.
* `tf.complex128`: 128-bit double-precision complex.
