Skip to content

Commit 655b266

Browse files
Akshay Moditensorflower-gardener
Akshay Modi
authored and
tensorflower-gardener
committed
Apply "Raise exception in SWIG on bad TF_Status" to base.i
Minor fixes to make this work. PiperOrigin-RevId: 191457070
1 parent cfc886a commit 655b266

8 files changed

Lines changed: 54 additions & 76 deletions

File tree

‎tensorflow/python/client/tf_session.i‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ limitations under the License.
2323
#include "tensorflow/core/lib/strings/stringprintf.h"
2424
#include "tensorflow/core/public/version.h"
2525
#include "tensorflow/python/client/tf_session_helper.h"
26-
#include "tensorflow/python/lib/core/py_exception_registry.h"
2726

2827
// Helper function to convert a Python list of Tensors to a C++ vector of
2928
// TF_Outputs.
@@ -353,27 +352,6 @@ TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper{
353352
reinterpret_cast<const char*>($1.data), $1.length);
354353
}
355354

356-
// Typemaps to automatically raise a Python exception from bad output TF_Status.
357-
// TODO(b/77295559): expand this to all TF_Status* output params and deprecate
358-
// raise_exception_on_not_ok_status (currently it only affects the C API).
359-
%typemap(in, numinputs=0) TF_Status* status (TF_Status* status) {
360-
status = TF_NewStatus();
361-
$1 = status;
362-
}
363-
364-
%typemap(argout) TF_Status* status {
365-
TF_Code code = TF_GetCode($1);
366-
if (code != TF_OK) {
367-
PyObject* exc = tensorflow::PyExceptionRegistry::Lookup(code);
368-
// Arguments to OpError.
369-
PyObject* exc_args = Py_BuildValue("sss", nullptr, nullptr, TF_Message($1));
370-
TF_DeleteStatus($1);
371-
SWIG_SetErrorObj(exc, exc_args);
372-
SWIG_fail;
373-
}
374-
TF_DeleteStatus($1);
375-
}
376-
377355
// Converts input Python list of wrapped TF_Outputs into a single array
378356
%typemap(in) (const TF_Output* inputs, int num_inputs)
379357
(std::vector<TF_Output> inputs) {
@@ -784,7 +762,3 @@ def TF_Reset(target, containers=None, config=None):
784762
%include "tensorflow/python/client/tf_session_helper.h"
785763

786764
%unignoreall
787-
788-
// Clear "TF_Status* status" typemap so it doesn't affect other modules and
789-
// unexpectedly remove the TF_Status* argument from wrappers.
790-
%clear TF_Status* status;

‎tensorflow/python/eager/backprop.py‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
from tensorflow.python.eager import tape
3232
from tensorflow.python.framework import constant_op
3333
from tensorflow.python.framework import dtypes
34-
from tensorflow.python.framework import errors
3534
from tensorflow.python.framework import ops
3635
from tensorflow.python.framework import tensor_shape
3736
from tensorflow.python.ops import array_ops
@@ -50,12 +49,10 @@ def op_attr_type(op_type, attr_name):
5049
try:
5150
return _op_attr_type_cache[(op_type, attr_name)]
5251
except KeyError:
53-
with errors.raise_exception_on_not_ok_status() as status:
54-
h = context.context()._handle # pylint: disable=protected-access
55-
attr_type = pywrap_tensorflow.TFE_OpNameGetAttrType(
56-
h, op_type, attr_name, status)
57-
_op_attr_type_cache[(op_type, attr_name)] = attr_type
58-
return attr_type
52+
h = context.context()._handle # pylint: disable=protected-access
53+
attr_type = pywrap_tensorflow.TFE_OpNameGetAttrType(h, op_type, attr_name)
54+
_op_attr_type_cache[(op_type, attr_name)] = attr_type
55+
return attr_type
5956

6057

6158
def make_attr(attr_type, value):

‎tensorflow/python/eager/context.py‎

Lines changed: 22 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
from tensorflow.python import pywrap_tensorflow
2929
from tensorflow.python.framework import c_api_util
3030
from tensorflow.python.framework import device as pydev
31-
from tensorflow.python.framework import errors
3231
from tensorflow.python.util import compat
3332
from tensorflow.python.util import is_in_graph_mode
3433
from tensorflow.python.util import tf_contextlib
@@ -224,24 +223,21 @@ def _initialize_handle_and_devices(self):
224223
assert self._context_devices is None
225224
opts = pywrap_tensorflow.TFE_NewContextOptions()
226225
try:
227-
with errors.raise_exception_on_not_ok_status() as status:
228-
if self._config is not None:
229-
config_str = self._config.SerializeToString()
230-
pywrap_tensorflow.TFE_ContextOptionsSetConfig(
231-
opts, config_str, len(config_str), status)
232-
if self._device_policy is not None:
233-
pywrap_tensorflow.TFE_ContextOptionsSetDevicePlacementPolicy(
234-
opts, self._device_policy)
235-
if self._execution_mode == ASYNC:
236-
pywrap_tensorflow.TFE_ContextOptionsSetAsync(opts, True)
237-
self._context_handle = pywrap_tensorflow.TFE_NewContext(opts, status)
226+
if self._config is not None:
227+
config_str = self._config.SerializeToString()
228+
pywrap_tensorflow.TFE_ContextOptionsSetConfig(opts, config_str)
229+
if self._device_policy is not None:
230+
pywrap_tensorflow.TFE_ContextOptionsSetDevicePlacementPolicy(
231+
opts, self._device_policy)
232+
if self._execution_mode == ASYNC:
233+
pywrap_tensorflow.TFE_ContextOptionsSetAsync(opts, True)
234+
self._context_handle = pywrap_tensorflow.TFE_NewContext(opts)
238235
finally:
239236
pywrap_tensorflow.TFE_DeleteContextOptions(opts)
240237
# Store list of devices
241238
self._context_devices = []
242-
with errors.raise_exception_on_not_ok_status() as status:
243-
device_list = pywrap_tensorflow.TFE_ContextListDevices(
244-
self._context_handle, status)
239+
device_list = pywrap_tensorflow.TFE_ContextListDevices(
240+
self._context_handle)
245241
try:
246242
self._num_gpus = 0
247243
for i in range(pywrap_tensorflow.TF_DeviceListCount(device_list)):
@@ -412,9 +408,7 @@ def set_execution_mode(self, mode):
412408
if mode is None:
413409
mode = SYNC
414410
self._eager_context.execution_mode = mode
415-
with errors.raise_exception_on_not_ok_status() as status:
416-
pywrap_tensorflow.TFE_ContextSetAsyncForThread(self._handle,
417-
mode == ASYNC, status)
411+
pywrap_tensorflow.TFE_ContextSetAsyncForThread(self._handle, mode == ASYNC)
418412

419413
@tf_contextlib.contextmanager
420414
def execution_mode(self, mode):
@@ -428,8 +422,7 @@ def execution_mode(self, mode):
428422

429423
def async_wait(self):
430424
"""Waits for ops dispatched in ASYNC mode to finish."""
431-
with errors.raise_exception_on_not_ok_status() as status:
432-
pywrap_tensorflow.TFE_ContextAsyncWait(self._handle, status)
425+
pywrap_tensorflow.TFE_ContextAsyncWait(self._handle)
433426

434427
def async_clear_error(self):
435428
"""Clears errors raised during ASYNC execution."""
@@ -449,11 +442,9 @@ def add_function(self, fn):
449442
Args:
450443
fn: A wrapped TF_Function (returned from TF_GraphToFunction_wrapper).
451444
"""
452-
with errors.raise_exception_on_not_ok_status() as status:
453-
pywrap_tensorflow.TFE_ContextAddFunction(
454-
self._handle, # pylint: disable=protected-access
455-
fn,
456-
status)
445+
pywrap_tensorflow.TFE_ContextAddFunction(
446+
self._handle, # pylint: disable=protected-access
447+
fn)
457448

458449
def add_function_def(self, fdef):
459450
"""Add a function definition to the context.
@@ -465,12 +456,10 @@ def add_function_def(self, fdef):
465456
fdef: A FunctionDef protocol buffer message.
466457
"""
467458
fdef_string = fdef.SerializeToString()
468-
with errors.raise_exception_on_not_ok_status() as status:
469-
pywrap_tensorflow.TFE_ContextAddFunctionDef(
470-
self._handle, # pylint: disable=protected-access
471-
fdef_string,
472-
len(fdef_string),
473-
status)
459+
pywrap_tensorflow.TFE_ContextAddFunctionDef(
460+
self._handle, # pylint: disable=protected-access
461+
fdef_string,
462+
len(fdef_string))
474463

475464
def add_post_execution_callback(self, callback):
476465
"""Add a post-execution callback to the context.
@@ -545,9 +534,8 @@ def export_run_metadata(self):
545534
if not self._context_handle:
546535
return None
547536
with c_api_util.tf_buffer() as buffer_:
548-
with errors.raise_exception_on_not_ok_status() as status:
549-
pywrap_tensorflow.TFE_ContextExportRunMetadata(
550-
self._context_handle, buffer_, status)
537+
pywrap_tensorflow.TFE_ContextExportRunMetadata(
538+
self._context_handle, buffer_)
551539
proto_data = pywrap_tensorflow.TF_GetBuffer(buffer_)
552540
run_metadata = config_pb2.RunMetadata()
553541
run_metadata.ParseFromString(compat.as_bytes(proto_data))

‎tensorflow/python/eager/imperative_grad.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import collections
2222

2323
from tensorflow.python import pywrap_tensorflow
24-
from tensorflow.python.framework import errors
2524

2625

2726
VSpace = collections.namedtuple(
@@ -60,6 +59,5 @@ def imperative_grad(
6059
or if only non-differentiable functions of the source were used in the
6160
computation of target.
6261
"""
63-
with errors.raise_exception_on_not_ok_status() as status:
64-
return pywrap_tensorflow.TFE_Py_TapeGradient(
65-
tape._tape, vspace, target, sources, output_gradients, status) # pylint: disable=protected-access
62+
return pywrap_tensorflow.TFE_Py_TapeGradient(
63+
tape._tape, vspace, target, sources, output_gradients) # pylint: disable=protected-access

‎tensorflow/python/grappler/item.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,7 @@ def __init__(self,
5151
self._BuildTFItem()
5252

5353
def IdentifyImportantOps(self, sort_topologically=False):
54-
with errors.raise_exception_on_not_ok_status() as status:
55-
return tf_item.TF_IdentifyImportantOps(self.tf_item, sort_topologically,
56-
status)
54+
return tf_item.TF_IdentifyImportantOps(self.tf_item, sort_topologically)
5755

5856
def GetOpProperties(self):
5957
ret_from_swig = tf_item.TF_GetOpProperties(self.tf_item)

‎tensorflow/python/lib/io/tf_record.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,7 @@ def tf_record_iterator(path, options=None):
7878
try:
7979
while True:
8080
try:
81-
with errors.raise_exception_on_not_ok_status() as status:
82-
reader.GetNext(status)
81+
reader.GetNext()
8382
except errors.OutOfRangeError:
8483
break
8584
yield reader.record()

‎tensorflow/python/platform/base.i‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,3 +229,25 @@ _COPY_TYPEMAPS(unsigned int, mode_t);
229229
%define final %enddef
230230
%define override %enddef
231231
#endif
232+
233+
// Typemaps to automatically raise a Python exception from bad output TF_Status.
234+
// TODO(b/77295559): expand this to all TF_Status* output params and deprecate
235+
// raise_exception_on_not_ok_status (currently it only affects the C API).
236+
%typemap(in, numinputs=0) TF_Status* status (TF_Status* status) {
237+
$1 = TF_NewStatus();
238+
}
239+
240+
%typemap(freearg) (TF_Status* status) {
241+
TF_DeleteStatus($1);
242+
}
243+
244+
%typemap(argout) TF_Status* status {
245+
TF_Code code = TF_GetCode($1);
246+
if (code != TF_OK) {
247+
PyObject* exc = tensorflow::PyExceptionRegistry::Lookup(code);
248+
// Arguments to OpError.
249+
PyObject* exc_args = Py_BuildValue("sss", nullptr, nullptr, TF_Message($1));
250+
SWIG_SetErrorObj(exc, exc_args);
251+
SWIG_fail;
252+
}
253+
}

‎tensorflow/python/pywrap_tfe.i‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16+
%include "tensorflow/python/platform/base.i"
17+
1618
%ignore "";
1719

1820
%rename("%s") TFE_NewContext;

0 commit comments

Comments
 (0)