TensorListElementShape should return scalar with value -1 if shape has unknown

saxenasaurabh · tensorflower-gardener · commit e5a615ae85bc · 2018-10-23T15:33:07.000-07:00
rank.

PiperOrigin-RevId: 218418543
diff --git a/tensorflow/core/kernels/list_kernels.cc b/tensorflow/core/kernels/list_kernels.cc
@@ -274,13 +274,22 @@ class TensorListElementShape : public OpKernel {
                     "list. Saw: '",
                     c->input(0).scalar<Variant>()().DebugString(), "'"));
     Tensor* result;
-    OP_REQUIRES_OK(c, c->allocate_output(
-                          0, TensorShape{l->element_shape.dims()}, &result));
-    for (int i = 0; i < l->element_shape.dims(); ++i) {
+    if (l->element_shape.unknown_rank()) {
+      OP_REQUIRES_OK(c, c->allocate_output(0, TensorShape({}), &result));
       if (result->dtype() == DT_INT32) {
-        result->flat<int32>()(i) = l->element_shape.dim_size(i);
+        result->scalar<int32>()() = -1;
       } else {
-        result->flat<int64>()(i) = l->element_shape.dim_size(i);
+        result->scalar<int64>()() = -1;
+      }
+    } else {
+      OP_REQUIRES_OK(c, c->allocate_output(
+                            0, TensorShape{l->element_shape.dims()}, &result));
+      for (int i = 0; i < l->element_shape.dims(); ++i) {
+        if (result->dtype() == DT_INT32) {
+          result->flat<int32>()(i) = l->element_shape.dim_size(i);
+        } else {
+          result->flat<int64>()(i) = l->element_shape.dim_size(i);
+        }
       }
     }
   }
diff --git a/tensorflow/python/kernel_tests/list_ops_test.py b/tensorflow/python/kernel_tests/list_ops_test.py
@@ -678,6 +678,13 @@ def testZerosLikeVariant(self):
       self.assertAllEqual(
           self.evaluate(outputs[0]), np.zeros((2,), dtype=dtype.as_numpy_dtype))
 
+  @test_util.run_in_graph_and_eager_modes
+  def testElementShape(self):
+    l = list_ops.empty_tensor_list(
+        element_dtype=dtypes.float32, element_shape=-1)
+    shape = list_ops.tensor_list_element_shape(l, shape_type=dtypes.int32)
+    self.assertEqual(self.evaluate(shape), -1)
+
 
 if __name__ == "__main__":
   test.main()
