From e4d3e2b92c53acce0b4111224e6196698eb01c18 Mon Sep 17 00:00:00 2001 From: Jiayang Liu Date: Thu, 22 Oct 2026 11:20:07 +0800 Subject: [PATCH 1/2] Fix `tf.linalg.expm` for half precision inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `tf.linalg.expm` documents `float16` as a supported input dtype, but the implementation ends in a `MatrixSolve` call for which no CPU/GPU kernel is registered for `DT_HALF`, so eager execution fails with NotFoundError: Could not find device for node: {{node MatrixSolve}} = MatrixSolve[T=DT_HALF, adjoint=false] `bfloat16` fails even earlier, with `ValueError: tf.linalg.expm does not support matrices of type bfloat16`. Accumulate half precision inputs in `float32` and cast the result back to the input dtype. This keeps the documented dtype contract, matches how other ops handle missing half precision kernels, and does not change the numerics for `float32`/`float64`/`complex` inputs. The Padé coefficients selected below are tuned for at least `float32` precision, so upcasting is also more accurate than running the same recurrence in 16 bits. The half precision test cases compare against a `float64` reference and use a `1e-2` tolerance (the existing `1e-3` tolerance for the other dtypes is unchanged). Fixes #121912 --- .../linalg/matrix_exponential_op_test.py | 22 ++++++++++++++----- tensorflow/python/ops/linalg/linalg_impl.py | 22 +++++++++++++++---- 2 files changed, 34 insertions(+), 10 deletions(-) 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..c1955cb53ce 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,36 @@ 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..f538995d5d3 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,15 @@ 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:])) + else: + result = array_ops.reshape(result, + batch_shape.concatenate(result.shape[-2:])) + if original_dtype != matrix.dtype: + result = math_ops.cast(result, original_dtype) + return result @tf_export('linalg.banded_triangular_solve', v1=[]) From ca9a6ac8f5751449d0226956b0f7934cbd898295 Mon Sep 17 00:00:00 2001 From: Jiayang Liu Date: Thu, 22 Oct 2026 12:05:11 +0800 Subject: [PATCH 2/2] Check for half precision explicitly when casting the expm result back Review feedback: comparing `original_dtype != matrix.dtype` requires tracing back to where `matrix` was promoted to know why the dtypes differ. Testing for `float16`/`bfloat16` directly mirrors the upcast guard at the top of the function. No behaviour change; `matrix_exponential_op_test.py` still reports 76 tests passing. --- tensorflow/python/ops/linalg/linalg_impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/linalg/linalg_impl.py b/tensorflow/python/ops/linalg/linalg_impl.py index f538995d5d3..161884b7310 100644 --- a/tensorflow/python/ops/linalg/linalg_impl.py +++ b/tensorflow/python/ops/linalg/linalg_impl.py @@ -395,7 +395,7 @@ def matrix_exponential(input, name=None): # pylint: disable=redefined-builtin else: result = array_ops.reshape(result, batch_shape.concatenate(result.shape[-2:])) - if original_dtype != matrix.dtype: + if original_dtype in (dtypes.float16, dtypes.bfloat16): result = math_ops.cast(result, original_dtype) return result