Skip to content
Permalink
Browse files Browse the repository at this point in the history
Validate proto.dtype() before calling set_dtype().
This prevents a `DCHECK`-fail when the proto contains an invalid dtype for a tensor shape with 0 elements or for an incomplete tensor shape.

PiperOrigin-RevId: 408369083
Change-Id: Ia21a3e3d62a90d642a4561f08f3b543e5ad00c46
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 8, 2021
1 parent 9f3eb61 commit 5b491cd
Showing 1 changed file with 9 additions and 0 deletions.
9 changes: 9 additions & 0 deletions tensorflow/core/framework/tensor.cc
Expand Up @@ -983,6 +983,15 @@ bool Tensor::FromProto(Allocator* a, const TensorProto& proto) {
dtype_error = true, dtype_error = true);
}
if (dtype_error || p == nullptr) return false;
} else {
// Handle the case of empty tensors (N = 0) or tensors with incomplete shape
// (N = -1). All other values of `shape.num_elements()` should be invalid by
// construction.
// Here, we just need to validate that the `proto.dtype()` value is valid.
bool dtype_error = false;
CASES_WITH_DEFAULT(proto.dtype(), break, dtype_error = true,
dtype_error = true);
if (dtype_error) return false;
}
shape_ = shape;
set_dtype(proto.dtype());
Expand Down

0 comments on commit 5b491cd

Please sign in to comment.