Check that the type of an implicitly dereferenced tensor matches the expected input...

The dtype of a tensor reference can change between the point when it is "produced" by an
operation and consumed by the next operation. This evades checks in the executor that the
type of tensor on each edge matches the type signatures of the producing and consuming operation, which could lead to undefined behavior. Although there is no existing operation that changes the type of a tensor reference, it is possible to use the OpKernelContext API to do so, so we add a further check in the runtime to defend against operations that might be added in the future.

PiperOrigin-RevId: 184356242

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index f515590..e3416da 100644 (file)
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -1776,6 +1776,19 @@ Status ExecutorState::PrepareInputs(const NodeItem& item, Entry* first_input,
         entry->ref_mu = nullptr;
 
         inp->tensor = entry->val.get();
+        // The dtype of entry->ref could have been changed by another operation
+        // that ran after the operation that "produced" it executed, so
+        // re-validate that the type of the dereferenced tensor matches the
+        // expected input type.
+        if (item.input_type(i) != inp->tensor->dtype()) {
+          return AttachDef(
+              errors::InvalidArgument(
+                  i, "-th input expects type ",
+                  DataTypeString(item.input_type(i)),
+                  " but automatically dereferenced input tensor has type ",
+                  DataTypeString(inp->tensor->dtype())),
+              item.kernel->def());
+        }
       }
     }
   }