Add stricter type checking for tf.math.real

tensorflow/core/framework/types.h

@@ -469,6 +469,11 @@ inline bool DataTypeIsFloating(DataType dt) {
   return kDataTypeIsFloating.Contains(dt);
 }
 
+// Returns true iff 'dt' is a numeric type.
+inline bool DataTypeIsNumeric(DataType dt) {
+  return kNumberTypes.Contains(dt);
+}
+
 // Returns true iff 'dt' is a complex type.
 constexpr DataTypeSet kDataTypeIsComplex =
     ToSet(DT_COMPLEX64) | ToSet(DT_COMPLEX128);

tensorflow/python/framework/dtypes.cc

@@ -120,6 +120,12 @@ PYBIND11_MODULE(_dtypes, m) {
             return tensorflow::BaseType(self) == tensorflow::DT_BOOL;
           },
           "Returns whether this is a boolean data type.")
+      .def_property_readonly(
+          "is_numeric",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsNumeric(tensorflow::BaseType(self));
+          },
+          "Returns whether this is a numeric data type.")
       .def_property_readonly(
           "is_complex",
           [](tensorflow::DataType self) {

tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py

@@ -1327,6 +1327,31 @@ def test_coeffs_raise(self):
     coeffs = {}
     with self.assertRaisesRegex(ValueError, "Argument coeffs must be list"):
       math_ops.polyval(coeffs, x)
+
+class RealTest(test.TestCase):
+
+  def _run_test(self, input_values, expected_values):
+    res = real(input_values)
+    self.assertListEqual(res, expected_values)
+
+  def test_real(self):
+    test_cases = [
+      # Complex tensor
+      (np.complex64, [-2.25 + 4.75j, 3.25 + 5.75j], [-2.25, 3.25]),
+      (np.complex128, [-2.25 + 4.75j, 3.25 + 5.75j], [-2.25, 3.25]),
+      # Real tensor
+      (np.float32, [1.0, 2.0, 3.0], [1.0, 2.0, 3.0]),
+      (np.float64, [1.0, 2.0, 3.0], [1.0, 2.0, 3.0]),
+    ]
+
+    for dtype, input_values, expected_values in test_cases:
+      with self.subTest(dtype=dtype):
+        self._run_test(input_values, expected_values)
+
+  def test_real_raises_error_for_non_numeric_tensor(self):
+    x = np.array(["Hello", "World"])
+    with self.assertRaisesRegex(TypeError, "input must be a numeric tensor"):
+      self._run_test(x, None)
 
 
 if __name__ == "__main__":

tensorflow/python/ops/math_ops.py

@@ -820,8 +820,11 @@ def real(input, name=None):
     if input.dtype.is_complex:
       real_dtype = input.dtype.real_dtype
       return gen_math_ops.real(input, Tout=real_dtype, name=name)
-    else:
+    elif input.dtype.is_numeric:
       return input
+    else:
+      raise TypeError("input must be a numeric tensor, but got "
+                      "tensor with dtype {}".format(input.dtype))
 
 
 @tf_export("math.imag", v1=["math.imag", "imag"])
@@ -1048,7 +1051,8 @@ def saturate_cast(value, dtype, name=None):
         # Clamp real and imag components separately, if required.
         real_in_dtype = in_dtype.real_dtype
         real_out_dtype = dtype.real_dtype
-        if real_in_dtype.min < real_out_dtype.min or real_in_dtype.max > real_out_dtype.max:
+        if (real_in_dtype.min < real_out_dtype.min or 
+            real_in_dtype.max > real_out_dtype.max):
           value = gen_math_ops._clip_by_value(
               value,
               ops.convert_to_tensor(