[TF2XLA] Attach Python stack trace to all INVALID_ARGUMENT Status objects created from the bridge

PiperOrigin-RevId: 356400325
Change-Id: I1b8fc2a771bcb6492ccde5fbeb14844154b38791
This commit is contained in:
George Karpenkov 2021-02-08 18:24:06 -08:00 committed by TensorFlower Gardener
parent dac302eeed
commit 2c596172d2
2 changed files with 23 additions and 22 deletions

View File

@ -154,8 +154,7 @@ class TensorListReserveOp : public XlaOpKernel {
"XLA compilation requires a fixed tensor list size. Set the number "
"of elements. This could also happen if you're using a TensorArray "
"in a while loop that does not have its maximum_iteration set, you "
"can fix this by setting maximum_iteration to a suitable value.",
ctx->StackTrace()));
"can fix this by setting maximum_iteration to a suitable value."));
// If element shape is compile time constant and it's not "unknown rank"
// shape (-1), create an initialized TensorList. Otherwise create an
@ -225,8 +224,7 @@ class EmptyTensorListOp : public XlaOpKernel {
"the max number of elements. This could also happen if "
"you're using a TensorArray in a while loop that does not "
"have its maximum_iteration set, you can fix this by "
"setting maximum_iteration to a suitable value.",
ctx->StackTrace()));
"setting maximum_iteration to a suitable value."));
if (dtype_ != DT_VARIANT) {
// We are creating a non-nested TensorList.
@ -294,8 +292,7 @@ class TensorListElementShapeOp : public XlaOpKernel {
OP_REQUIRES_OK(ctx,
(IsTensorListInitialized(ctx->Input(0), &is_initialized)));
OP_REQUIRES(ctx, is_initialized,
errors::InvalidArgument("TensorList is not initialized",
ctx->StackTrace()));
errors::InvalidArgument("TensorList is not initialized"));
// Only non-nested TensorList is supported for now.
bool is_nested;
@ -351,8 +348,7 @@ class TensorListGetItemOp : public XlaOpKernel {
OP_REQUIRES_OK(ctx,
(IsTensorListInitialized(ctx->Input(0), &is_initialized)));
OP_REQUIRES(ctx, is_initialized,
errors::InvalidArgument("TensorList is not initialized",
ctx->StackTrace()));
errors::InvalidArgument("TensorList is not initialized"));
// Only non-nested TensorList is supported for now.
bool is_nested;
@ -390,8 +386,7 @@ class TensorListGatherOp : public XlaOpKernel {
OP_REQUIRES_OK(ctx,
(IsTensorListInitialized(ctx->Input(0), &is_initialized)));
OP_REQUIRES(ctx, is_initialized,
errors::InvalidArgument("TensorList is not initialized",
ctx->StackTrace()));
errors::InvalidArgument("TensorList is not initialized"));
// Only non-nested TensorList is supported for now.
bool is_nested;
@ -442,8 +437,7 @@ class TensorListStackOp : public XlaOpKernel {
OP_REQUIRES_OK(ctx,
(IsTensorListInitialized(ctx->Input(0), &is_initialized)));
OP_REQUIRES(ctx, is_initialized,
errors::InvalidArgument("TensorList is not initialized",
ctx->StackTrace()));
errors::InvalidArgument("TensorList is not initialized"));
// Only non-nested TensorList is supported for now.
bool is_nested;
@ -474,8 +468,7 @@ class TensorListConcatOp : public XlaOpKernel {
bool is_initialized;
OP_REQUIRES_OK(ctx, (IsTensorListInitialized(input, &is_initialized)));
OP_REQUIRES(ctx, is_initialized,
errors::InvalidArgument("TensorList is not initialized",
ctx->StackTrace()));
errors::InvalidArgument("TensorList is not initialized"));
// Only non-nested TensorList is supported for now.
bool is_nested;
@ -673,8 +666,7 @@ class TensorListPopBackOp : public XlaOpKernel {
OP_REQUIRES_OK(ctx,
(IsTensorListInitialized(ctx->Input(0), &is_initialized)));
OP_REQUIRES(ctx, is_initialized,
errors::InvalidArgument("TensorList is not initialized",
ctx->StackTrace()));
errors::InvalidArgument("TensorList is not initialized"));
xla::XlaOp list = ctx->Input(0);
xla::XlaOp list_result, element_result;

View File

@ -177,8 +177,7 @@ Status XlaOpKernelContext::ConstantInputReshaped(
"This error means that a shape or dimension argument could not be "
"evaluated at compile time, usually because the value of the argument "
"depends on a parameter to the computation, on a variable, or on a "
"stateful operation such as a random number generator.",
StackTrace());
"stateful operation such as a random number generator.");
}
Tensor temp(constant->dtype());
@ -664,19 +663,29 @@ Status XlaOpKernelContext::AssignVariable(absl::string_view name, DataType type,
builder());
}
static Status GetStatusWithStackTrace(const Status& s,
const XlaOpKernelContext* ctx) {
if (s.code() == error::INVALID_ARGUMENT) {
return Status{s.code(),
absl::StrCat(s.error_message(), "\n", ctx->StackTrace())};
}
return s;
}
void XlaOpKernelContext::CtxFailure(const Status& s) {
context_->CtxFailure(s);
context_->CtxFailure(GetStatusWithStackTrace(s, this));
}
void XlaOpKernelContext::CtxFailureWithWarning(const Status& s) {
context_->CtxFailureWithWarning(s);
context_->CtxFailureWithWarning(GetStatusWithStackTrace(s, this));
}
void XlaOpKernelContext::CtxFailure(const char* file, int line,
const Status& s) {
context_->CtxFailure(file, line, s);
context_->CtxFailure(file, line, GetStatusWithStackTrace(s, this));
}
void XlaOpKernelContext::CtxFailureWithWarning(const char* file, int line,
const Status& s) {
context_->CtxFailureWithWarning(file, line, s);
context_->CtxFailureWithWarning(file, line, GetStatusWithStackTrace(s, this));
}
const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMax(