experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Raise error type corresponding to status code instead of generic RuntimeError.

Browse Source 
PiperOrigin-RevId: 317367352
Change-Id: I35378b88a33269ac225632ae848398b819c694a1
This commit is contained in:
Haoyu Zhang 2020-06-19 13:19:39 -07:00 committed by TensorFlower Gardener
parent c575e2ba93
commit 72d30dfb8b
1 changed files with 21 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
tensorflow/python/eager/pywrap_tensor.cc
View File

@ -180,6 +180,15 @@ int ConvertDeviceName(PyObject* obj, const char** dst) {
return 1;
}
void RaiseExceptionTypeFromTFStatus(TF_Status* status) {
TF_Code code = TF_GetCode(status);
PyObject* exception = tensorflow::PyExceptionRegistry::Lookup(code);
PyErr_SetObject(exception,
pybind11::make_tuple(pybind11::none(), pybind11::none(),
TF_Message(status))
.ptr());
}
} // namespace
namespace tensorflow {
@ -305,13 +314,7 @@ TFE_TensorHandle* ConvertToEagerTensorUncached(TFE_Context* ctx,
device_name, status.get()));
const TF_Code code = TF_GetCode(status.get());
if (code != TF_OK) {
// Instead of raising a generic RuntimeError, raise an exception type
// based on the status error code.
PyObject* exception = PyExceptionRegistry::Lookup(code);
PyErr_SetObject(exception,
pybind11::make_tuple(pybind11::none(), pybind11::none(),
TF_Message(status.get()))
.ptr());
RaiseExceptionTypeFromTFStatus(status.get());
return nullptr;
}
}
@ -512,7 +515,9 @@ static PyObject* EagerTensor_datatype_enum(EagerTensor* self) {
static PyObject* EagerTensor_shape_tuple(EagerTensor* self) {
auto handle = self->handle;
int n = TFE_TensorHandleNumDims(handle, &self->status);
if (MaybeRaiseExceptionFromTFStatus(&self->status, nullptr)) {
TF_Code code = TF_GetCode(&self->status);
if (code != TF_OK) {
RaiseExceptionTypeFromTFStatus(&self->status);
// Cleanup self->status before returning.
self->status.status = tensorflow::Status::OK();
return nullptr;
@ -522,13 +527,18 @@ static PyObject* EagerTensor_shape_tuple(EagerTensor* self) {
for (int i = 0; i < n; ++i) {
PyObject* dim =
PyLong_FromLongLong(TFE_TensorHandleDim(handle, i, &self->status));
if (MaybeRaiseExceptionFromTFStatus(&self->status, nullptr) ||
dim == nullptr || PyTuple_SetItem(shape, i, dim) != 0) {
code = TF_GetCode(&self->status);
if (code != TF_OK || dim == nullptr ||
PyTuple_SetItem(shape, i, dim) != 0) {
if (code != TF_OK) {
RaiseExceptionTypeFromTFStatus(&self->status);
} else {
PyErr_SetString(PyExc_RuntimeError, "Error while creating shape");
}
// Cleanup self->status before returning.
self->status.status = tensorflow::Status::OK();
Py_DECREF(shape);
if (dim != nullptr) Py_DECREF(dim);
PyErr_SetString(PyExc_RuntimeError, "Error while creating shape");
return nullptr;
}
}