diff options
| author |  Asim Shankar <ashankar@google.com> | 2018-06-15 16:26:35 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-06-15 16:31:01 -0700 |
| commit | 4d8a66c5b29428b709f4f54b566a44902ea8173e (patch) | |
| tree | 717b000e9284c2c9ae0710c32df14bc36f86dda7 /tensorflow/python/BUILD | |
| parent | 96100f90a90bb2db905f50617cbb5e7928480667 (diff) | |
[py_func]: Fix #20021
* EagerPyFunc now validates its assumption that returned tensors are backed by memory on the same device that the EagerPyFunc kernel executed on.
* Make the Python trampolining mechanism ensure that this requirement of the kernel is met.
* Allow tf.contrib.eager.py_func to execute correctly on devices other than CPU and GPU:0.
Prior to this change, tf.contrib.eager.py_func() would copy data from CPU to GPU:0 if necessary, but not the other way around. As a result, the assumptions made by the EagerPyFunc kernel implementation about the placement of returned tensors would be violated.
The test added in py_func_test.py, when executed on a machine with a GPU will:
- Fail with a segmentation fault (dereferencing GPU memory) without the changes to py_func.cc and script_ops.py
- Fail with an error message with the change to py_func.cc but without the change to script_ops.py
- Pass with changes to py_func.cc and script_ops.py
PiperOrigin-RevId: 200792057
Diffstat (limited to 'tensorflow/python/BUILD')
| -rw-r--r-- | tensorflow/python/BUILD | 1 |
1 files changed, 1 insertions, 0 deletions
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index 39e0cafd93..f3a848b7df 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -2428,6 +2428,7 @@ py_library( srcs = ["ops/script_ops.py"], srcs_version = "PY2AND3", deps = [ + ":array_ops", ":framework_for_generated_wrappers", ":script_ops_gen", "//third_party/py/numpy", |