Add complex128 support for tf.as_string

The tf.as_string supports most of the data types (including complex64) but not complex128. This fix add complex128 support for tf.as_string. Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
tensorflow · Jul 31, 2018 · 52bb2b0 · 52bb2b0
1 parent 2826d12
commit 52bb2b0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/tensorflow/core/kernels/as_string_op.cc b/tensorflow/core/kernels/as_string_op.cc
@@ -47,6 +47,7 @@ class AsStringOp : public OpKernel {
       case DT_FLOAT:
       case DT_DOUBLE:
       case DT_COMPLEX64:
+      case DT_COMPLEX128:
         break;
       default:
         OP_REQUIRES(ctx, !(scientific || shortest),
@@ -83,6 +84,7 @@ class AsStringOp : public OpKernel {
       case DT_FLOAT:
       case DT_DOUBLE:
       case DT_COMPLEX64:
+      case DT_COMPLEX128:
         if (shortest) {
           strings::Appendf(&format_, "g");
         } else if (scientific) {
@@ -100,7 +102,7 @@ class AsStringOp : public OpKernel {
                                             DataTypeString(dtype)));
     }
 
-    if (dtype == DT_COMPLEX64) {
+    if (dtype == DT_COMPLEX64 || dtype == DT_COMPLEX128) {
       format_ = strings::Printf("(%s,%s)", format_.c_str(), format_.c_str());
     }
   }
@@ -144,6 +146,13 @@ class AsStringOp : public OpKernel {
               format_.c_str(), input_flat(i).real(), input_flat(i).imag());
         }
       } break;
+      case (DT_COMPLEX128): {
+        const auto& input_flat = input_tensor->flat<complex128>();
+        for (int i = 0; i < input_flat.size(); ++i) {
+          output_flat(i) = strings::Printf(
+              format_.c_str(), input_flat(i).real(), input_flat(i).imag());
+        }
+      } break;
       default:
         bool can_encode_type = false;
         OP_REQUIRES(context, can_encode_type,

diff --git a/tensorflow/core/ops/string_ops.cc b/tensorflow/core/ops/string_ops.cc
@@ -78,7 +78,9 @@ REGISTER_OP("ReduceJoin")
 REGISTER_OP("AsString")
     .Input("input: T")
     .Output("output: string")
-    .Attr("T: {int8, int16, int32, int64, complex64, float, double, bool}")
+    .Attr(
+        "T: {int8, int16, int32, int64, complex64, complex128, float, double, "
+        "bool}")
     .Attr("precision: int = -1")
     .Attr("scientific: bool = false")
     .Attr("shortest: bool = false")