Internal change
PiperOrigin-RevId: 293241641
Change-Id: I93ab396e90f8aa70367869128a104bcf87c3f641
diff --git a/tensorflow/compiler/jit/compilability_check_util.cc b/tensorflow/compiler/jit/compilability_check_util.cc
index 14ade0e..8a5f682 100644
--- a/tensorflow/compiler/jit/compilability_check_util.cc
+++ b/tensorflow/compiler/jit/compilability_check_util.cc
@@ -266,9 +266,9 @@
s = lib_runtime->Instantiate(function.name(), AttrSlice(&function.attr()),
&handle);
}
-
if (!s.ok()) {
- std::string uncompilable_reason = "could not instantiate call";
+ std::string uncompilable_reason =
+ absl::StrCat("could not instantiate call: '", function.name(), "'");
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting " << call_def.DebugString() << ": "
diff --git a/tensorflow/compiler/jit/xla_kernel_creator_util.cc b/tensorflow/compiler/jit/xla_kernel_creator_util.cc
index 167d351..0da97d8 100644
--- a/tensorflow/compiler/jit/xla_kernel_creator_util.cc
+++ b/tensorflow/compiler/jit/xla_kernel_creator_util.cc
@@ -143,11 +143,11 @@
}
string message = absl::StrCat(
"Function invoked by the following node is not compilable: ",
- node_def.ShortDebugString(), ".\n");
- absl::StrAppend(&message, "Uncompilable nodes:\n");
+ SummarizeNodeDef(node_def), ".\n");
+ absl::StrAppend(&message, "Uncompilable nodes:");
for (const auto& node_info : uncompilable_node_info) {
string node_message =
- absl::StrCat("\t", node_info.name, ": ",
+ absl::StrCat("\n", node_info.name, ": ",
node_info.uncompilable_reason, "\n", "\tStacktrace:\n");
for (const auto& stack_frame : node_info.stack_trace) {
absl::StrAppendFormat(&node_message, "\t\tNode: %s, function: %s\n",
@@ -156,7 +156,6 @@
absl::StrAppend(&message, node_message);
}
VLOG(1) << message;
- // node_def is calling a function that XLA can't compile.
return errors::InvalidArgument(message);
}
diff --git a/tensorflow/core/framework/op_kernel.cc b/tensorflow/core/framework/op_kernel.cc
index 1fe8c19..a87fbc9 100644
--- a/tensorflow/core/framework/op_kernel.cc
+++ b/tensorflow/core/framework/op_kernel.cc
@@ -24,6 +24,7 @@
#include <vector>
#include "absl/base/call_once.h"
+#include "absl/strings/match.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/attr_value_util.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
@@ -1389,6 +1390,7 @@
device_type, node_name, has_experimental_debug_info,
experimental_debug_info, node_op, node_attrs, ®, &was_attr_mismatch));
if (reg == nullptr) {
+ std::string device_str = DeviceTypeString(device_type);
Status s = errors::NotFound(
"No registered '", node_op, "' OpKernel for ",
DeviceTypeString(device_type), " devices compatible with node ",
@@ -1400,8 +1402,14 @@
"Requested Attributes: ",
SummarizeAttrsHelper(node_attrs, node_device));
}
- errors::AppendToMessage(&s,
- ". Registered:", KernelsRegisteredForOp(node_op));
+
+ // Do not print kernel registrations for other devices when using _JIT
+ // devices for compilation.
+ if (!absl::StrContains(device_str, "JIT")) {
+ errors::AppendToMessage(
+ &s, ". Registered:", KernelsRegisteredForOp(node_op));
+ }
+
return s;
}
if (def != nullptr) *def = ®->def;
