Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Replace CHECK with returning an InternalError on failing to create py… #58576

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
12 changes: 8 additions & 4 deletions tensorflow/python/lib/core/py_func.cc
Expand Up @@ -83,8 +83,8 @@ bool IsCPUDevice(const Device* d) {
return d == nullptr || d->tensorflow_accelerator_device_info() == nullptr;
}

// Givens the 'call', prepares the token and inputs as a python tuple
// that is appropriate for calling the trampoline.
// Given the 'call', prepares the token and inputs as a python tuple that is
// appropriate for calling the trampoline.
Status MakeArgTuple(const PyCall* call, TFE_Context* ctx, PyObject** tuple) {
int64_t n = call->ins.size();
PyObject* lst = PyList_New(n);
Expand Down Expand Up @@ -119,8 +119,12 @@ Status MakeArgTuple(const PyCall* call, TFE_Context* ctx, PyObject** tuple) {
PyList_SetItem(lst, i, arg);
}
*tuple = Py_BuildValue("(ssN)", call->token.c_str(), device_name, lst);
CHECK(*tuple);
return Status::OK();
if (*tuple == nullptr) {
return errors::Internal(
"Failed to create python tuple. Please make sure `token` is a "
"well-formed UTF-8 string.");
}
return OkStatus();
}

bool IsSingleNone(PyObject* obj) {
Expand Down
28 changes: 28 additions & 0 deletions tensorflow/python/ops/script_ops_test.py
Expand Up @@ -16,8 +16,11 @@

from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.framework import constant_op
from tensorflow.python.ops import gen_script_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import script_ops
from tensorflow.python.ops.script_ops import numpy_function
from tensorflow.python.platform import test
Expand Down Expand Up @@ -87,5 +90,30 @@ def func_stateful(a, b):
2) # as stateful, func is guaranteed to execute twice


class PyFunctionTest(test.TestCase):

@test_util.run_in_graph_and_eager_modes
def test_variable_arguments(self):

def plus(a, b):
return a + b

v1 = resource_variable_ops.ResourceVariable(1)
self.evaluate(v1.initializer)

actual_result = script_ops.eager_py_func(plus, [v1, 2], dtypes.int32)
expect_result = constant_op.constant(3, dtypes.int32)
self.assertAllEqual(actual_result, expect_result)

@test_util.run_in_graph_and_eager_modes
def test_fail_on_non_utf8_token(self):
value = constant_op.constant(value=[1, 2])
token = b"\xb0"
data_type = [dtypes.int32]
with self.assertRaises((errors.InternalError, UnicodeDecodeError)):
self.evaluate(
gen_script_ops.py_func(input=[value], token=token, Tout=data_type))


if __name__ == "__main__":
test.main()