wehub-resource-sync
127 lines
4.1 KiB
Python
127 lines
4.1 KiB
Python
# Copyright 2025 The OpenXLA Authors.
|
|
#
|
|
# 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.
|
|
# ==============================================================================
|
|
import platform
|
|
|
|
import numpy as np
|
|
|
|
from tensorflow.compiler.tests import xla_test
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.platform import googletest
|
|
|
|
|
|
class CastTest(xla_test.XLATestCase):
|
|
|
|
def test_cast(self):
|
|
types = {
|
|
dtypes.bool,
|
|
dtypes.float32,
|
|
dtypes.float64,
|
|
dtypes.complex64,
|
|
dtypes.int32,
|
|
dtypes.int64,
|
|
dtypes.uint32,
|
|
dtypes.uint64,
|
|
}
|
|
with self.session() as session:
|
|
for src_type in types:
|
|
for dst_type in types:
|
|
self._test_cast(src_type, dst_type, session)
|
|
|
|
def test_cast_fp8(self):
|
|
if platform.system() == "Darwin":
|
|
# TODO(b/271327511): Fix issue where casts to FP8 very rarely result in
|
|
# NaN on Mac
|
|
self.skipTest("Casts to FP8 sometimes result in NaN on Mac")
|
|
fp8_types = {
|
|
dtypes.float8_e5m2,
|
|
dtypes.float8_e4m3fn,
|
|
dtypes.float8_e4m3fnuz,
|
|
dtypes.float8_e4m3b11fnuz,
|
|
dtypes.float8_e5m2fnuz,
|
|
}
|
|
other_types = {
|
|
dtypes.bool,
|
|
dtypes.float32,
|
|
dtypes.float64,
|
|
dtypes.complex64,
|
|
dtypes.int32,
|
|
dtypes.int64,
|
|
dtypes.uint32,
|
|
dtypes.uint64,
|
|
}
|
|
with self.session() as session:
|
|
for fp8_type in fp8_types:
|
|
for other_type in other_types | fp8_types:
|
|
self._test_cast(fp8_type, other_type, session)
|
|
self._test_cast(other_type, fp8_type, session)
|
|
|
|
def _test_cast(self, src_type, dst_type, session):
|
|
with self.subTest(src_type=src_type, dst_type=dst_type):
|
|
shapes = [[], [4], [2, 3], [2, 0, 4]]
|
|
src_np_dtype = src_type.as_numpy_dtype
|
|
dst_np_dtype = dst_type.as_numpy_dtype
|
|
|
|
for shape in shapes:
|
|
src = np.arange(np.prod(shape)).astype(src_np_dtype)
|
|
|
|
if src_type in self.complex_tf_types:
|
|
src += (np.arange(np.prod(shape)) * 2j).astype(src_np_dtype)
|
|
src = src.reshape(shape)
|
|
dst = src.astype(dst_np_dtype)
|
|
self.assert_op_output_matches_expected(
|
|
lambda x, dst_type=dst_type: math_ops.cast(x, dst_type),
|
|
src,
|
|
expected=dst,
|
|
local_session=session,
|
|
)
|
|
|
|
# Check special values.
|
|
if src_type.is_integer:
|
|
imin = np.iinfo(src_np_dtype).min
|
|
imax = np.iinfo(src_np_dtype).max
|
|
if src_type.is_unsigned:
|
|
src = np.array([imin, imax, 0, 1], dtype=src_np_dtype)
|
|
else:
|
|
src = np.array([imin, imax, 0, 1, -1], dtype=src_np_dtype)
|
|
elif src_type in self.float_tf_types:
|
|
if dst_type.is_integer:
|
|
imin = np.iinfo(dst_np_dtype).min
|
|
imax = np.iinfo(dst_np_dtype).max // 2
|
|
src = np.array([imin, imax, 0, 1], dtype=src_np_dtype)
|
|
elif dst_type in self.float_tf_types:
|
|
fmin = np.finfo(dst_np_dtype).min
|
|
fmax = np.finfo(dst_np_dtype).max
|
|
tiny = np.finfo(dst_np_dtype).tiny
|
|
eps = np.finfo(dst_np_dtype).eps
|
|
src = np.array(
|
|
[fmin, fmax, np.nan, eps, -eps, tiny, -tiny, np.inf, -np.inf],
|
|
dtype=src_np_dtype,
|
|
)
|
|
dst = src.astype(dst_np_dtype)
|
|
self.assert_op_output_matches_expected(
|
|
lambda x, dst_type=dst_type: math_ops.cast(x, dst_type),
|
|
src,
|
|
expected=dst,
|
|
local_session=session,
|
|
)
|
|
|
|
def test_give_me_a_name(self):
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
googletest.main()
|