diff --git a/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py b/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py index 534ed120638..94a656ce26d 100644 --- a/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py +++ b/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py @@ -22,6 +22,7 @@ from tensorflow.python.client import session from tensorflow.python.eager import backprop from tensorflow.python.eager import forwardprop from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops @@ -47,27 +48,40 @@ def np_expm(x): # pylint: disable=invalid-name return y +# Half precision dtypes are accumulated in float32 by `matrix_exponential`. +_HALF_TYPES = (np.float16, dtypes.bfloat16.as_numpy_dtype) + + @test_util.run_all_without_tensor_float_32("Avoid TF32-based matmuls.") class ExponentialOpTest(test.TestCase): def _verifyExponential(self, x, np_type): inp = x.astype(np_type) + # Accumulating the Taylor series in 16-bit would itself be inaccurate, so + # half precision inputs are compared against a float64 reference. + ref_type = np.float64 if np_type in _HALF_TYPES else np_type with test_util.use_gpu(): tf_ans = linalg_impl.matrix_exponential(inp) if x.size == 0: - np_ans = np.empty(x.shape, dtype=np_type) + np_ans = np.empty(x.shape, dtype=ref_type) else: if x.ndim > 2: - np_ans = np.zeros(inp.shape, dtype=np_type) + np_ans = np.zeros(inp.shape, dtype=ref_type) for i in itertools.product(*[range(x) for x in inp.shape[:-2]]): - np_ans[i] = np_expm(inp[i]) + np_ans[i] = np_expm(inp[i].astype(ref_type)) else: - np_ans = np_expm(inp) + np_ans = np_expm(inp.astype(ref_type)) out = self.evaluate(tf_ans) - self.assertAllClose(np_ans, out, rtol=1e-3, atol=1e-3) + tol = 1e-2 if np_type in _HALF_TYPES else 1e-3 + self.assertAllClose(np_ans, out, rtol=tol, atol=tol) def _verifyExponentialReal(self, x): - for np_type in [np.float32, np.float64]: + for np_type in [ + np.float16, + dtypes.bfloat16.as_numpy_dtype, + np.float32, + np.float64, + ]: self._verifyExponential(x, np_type) def _verifyExponentialComplex(self, x): diff --git a/tensorflow/python/ops/linalg/linalg_impl.py b/tensorflow/python/ops/linalg/linalg_impl.py index b46ca76e179..0e9ad3f6cc4 100644 --- a/tensorflow/python/ops/linalg/linalg_impl.py +++ b/tensorflow/python/ops/linalg/linalg_impl.py @@ -285,8 +285,10 @@ def matrix_exponential(input, name=None): # pylint: disable=redefined-builtin containing the exponential for all input submatrices `[..., :, :]`. Args: - input: A `Tensor`. Must be `float16`, `float32`, `float64`, `complex64`, or - `complex128` with shape `[..., M, M]`. + input: A `Tensor`. Must be `float16`, `bfloat16`, `float32`, `float64`, + `complex64`, or `complex128` with shape `[..., M, M]`. Half precision + inputs (`float16` and `bfloat16`) are accumulated in `float32` and cast + back to the input dtype. name: A name to give this `Op` (optional). Returns: @@ -303,6 +305,13 @@ def matrix_exponential(input, name=None): # pylint: disable=redefined-builtin matrix = ops.convert_to_tensor(input, name='input') if matrix.shape[-2:] == [0, 0]: return matrix + # There is no `MatrixSolve` kernel registered for `float16`/`bfloat16` on + # CPU or GPU (only XLA kernels exist), and the Pade coefficients selected + # below assume at least `float32` precision. Accumulate in `float32` and + # cast the result back, which is how other half-precision ops behave. + original_dtype = matrix.dtype + if original_dtype in (dtypes.float16, dtypes.bfloat16): + matrix = math_ops.cast(matrix, dtypes.float32) batch_shape = matrix.shape[:-2] if not batch_shape.is_fully_defined(): batch_shape = array_ops.shape(matrix)[:-2] @@ -380,10 +389,17 @@ def matrix_exponential(input, name=None): # pylint: disable=redefined-builtin _, result = while_loop.while_loop(c, b, [i, result]) if not matrix.shape.is_fully_defined(): - return array_ops.reshape( + result = array_ops.reshape( result, - array_ops.concat((batch_shape, array_ops.shape(result)[-2:]), axis=0)) - return array_ops.reshape(result, batch_shape.concatenate(result.shape[-2:])) + array_ops.concat((batch_shape, array_ops.shape(result)[-2:]), axis=0), + ) + else: + result = array_ops.reshape( + result, batch_shape.concatenate(result.shape[-2:]) + ) + if original_dtype in (dtypes.float16, dtypes.bfloat16): + result = math_ops.cast(result, original_dtype) + return result @tf_export('linalg.banded_triangular_solve', v1=[])