Merge pull request #129433 from Phoebus-Liu:fix/issue-121912-expm-16bit

PiperOrigin-RevId: 1005360375

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):

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=[])