Use PartialTensorShape instead of TensorShape.

mihaimaruseac · tensorflower-gardener · commit c2426bba00a0 · 2021-11-11T12:19:44.000-08:00
`TensorShape` constructor throws a CHECK-fail if shape is partial/overflows which the other doesn't. We are only determining the number of elements in the shape and partial shape should be used as it returns negative number when needed.

PiperOrigin-RevId: 409205384
Change-Id: Ia56542ff9ec758f2c9ffc7e4dcc9fa7eecd86e7b
diff --git a/tensorflow/core/framework/attr_value_util.cc b/tensorflow/core/framework/attr_value_util.cc
@@ -45,7 +45,7 @@ constexpr int kMaxTensorNestDepth = 100;
 // not fully defined return -1.
 int64_t TensorByteSize(const TensorProto& t) {
   // num_elements returns -1 if shape is not fully defined.
-  int64_t num_elems = TensorShape(t.tensor_shape()).num_elements();
+  int64_t num_elems = PartialTensorShape(t.tensor_shape()).num_elements();
   return num_elems < 0 ? -1 : num_elems * DataTypeSize(t.dtype());
 }
 

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ constexpr int kMaxTensorNestDepth = 100;`
`45`	`45`	`// not fully defined return -1.`
`46`	`46`	`int64_t TensorByteSize(const TensorProto& t) {`
`47`	`47`	`// num_elements returns -1 if shape is not fully defined.`
`48`		`- int64_t num_elems = TensorShape(t.tensor_shape()).num_elements();`
	`48`	`+ int64_t num_elems = PartialTensorShape(t.tensor_shape()).num_elements();`
`49`	`49`	`return num_elems < 0 ? -1 : num_elems * DataTypeSize(t.dtype());`
`50`	`50`	`}`
`51`	`51`