Use self->status in EagerTensor_numpy and EagerTensor_getbuffer

There is no reason to create a new status on each call.

PiperOrigin-RevId: 257594441
This commit is contained in:
Sergei Lebedev 2019-07-11 06:09:44 -07:00 committed by TensorFlower Gardener
parent 78ee8186db
commit b2b1c7181f

View File

@ -657,10 +657,11 @@ static PyObject* EagerTensor_copy_to_device(EagerTensor* self, PyObject* args,
// other. // other.
// Note that if `self` is not on CPU, we raise an Exception. // Note that if `self` is not on CPU, we raise an Exception.
static PyObject* EagerTensor_numpy(EagerTensor* self) { static PyObject* EagerTensor_numpy(EagerTensor* self) {
auto status = tensorflow::make_safe(TF_NewStatus()); auto* py_array = TFE_TensorHandleToNumpy(self->handle, self->status);
auto* py_array = TFE_TensorHandleToNumpy(self->handle, status.get()); if (MaybeRaiseExceptionFromTFStatus(self->status, PyExc_ValueError)) {
if (MaybeRaiseExceptionFromTFStatus(status.get(), PyExc_ValueError)) {
Py_XDECREF(py_array); Py_XDECREF(py_array);
// Cleanup self->status before returning.
TF_SetStatus(self->status, TF_OK, "");
return nullptr; return nullptr;
} else { } else {
return PyArray_Return(reinterpret_cast<PyArrayObject*>(py_array)); return PyArray_Return(reinterpret_cast<PyArrayObject*>(py_array));
@ -745,13 +746,14 @@ static int EagerTensor_getbuffer(EagerTensor* self, Py_buffer* view,
return -1; return -1;
} }
auto status = tensorflow::make_safe(TF_NewStatus());
// TensorHandleToNumpy is zero-copy for everything but DT_RESOURCE and // TensorHandleToNumpy is zero-copy for everything but DT_RESOURCE and
// DT_STRING so the following is only slightly slower than a NumPy-free // DT_STRING so the following is only slightly slower than a NumPy-free
// implementation. // implementation.
auto py_array = tensorflow::make_safe( auto py_array = tensorflow::make_safe(
TFE_TensorHandleToNumpy(self->handle, status.get())); TFE_TensorHandleToNumpy(self->handle, self->status));
if (MaybeRaiseExceptionFromTFStatus(status.get(), PyExc_BufferError)) { if (MaybeRaiseExceptionFromTFStatus(self->status, PyExc_BufferError)) {
// Cleanup self->status before returning.
TF_SetStatus(self->status, TF_OK, "");
return -1; return -1;
} }
if (PyObject_GetBuffer(py_array.get(), view, flags) < 0) { if (PyObject_GetBuffer(py_array.get(), view, flags) < 0) {