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

Arg/Ret ops should return/take tensors in device memory

Browse Source 
PiperOrigin-RevId: 238643502
This commit is contained in:
Igor Ganichev 2019-03-15 08:11:31 -07:00 committed by TensorFlower Gardener
parent 9a43dfeac5
commit ea6c5f5cc3
1 changed files with 3 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/compiler/jit/xla_device_ops.h
View File

@ -199,9 +199,7 @@ class XlaAssignVariableOp : public OpKernel {
Name("FIFOQueueV2").Device(DEVICE).HostMemory("handle"), FIFOQueueOp); \ Name("FIFOQueueV2").Device(DEVICE).HostMemory("handle"), FIFOQueueOp); \
\ \
REGISTER_KERNEL_BUILDER( \ REGISTER_KERNEL_BUILDER( \
Name(kArgOp).Device(DEVICE).HostMemory("output").TypeConstraint("T", \ Name(kArgOp).Device(DEVICE).TypeConstraint("T", TYPES), ArgOp); \
TYPES), \
ArgOp); \
REGISTER_KERNEL_BUILDER(Name(kArgOp) \ REGISTER_KERNEL_BUILDER(Name(kArgOp) \
.Device(DEVICE) \ .Device(DEVICE) \
.HostMemory("output") \ .HostMemory("output") \
@ -210,11 +208,8 @@ class XlaAssignVariableOp : public OpKernel {
REGISTER_KERNEL_BUILDER( \ REGISTER_KERNEL_BUILDER( \
Name(kArgOp).Device(DEVICE).TypeConstraint<Variant>("T"), ArgOp); \ Name(kArgOp).Device(DEVICE).TypeConstraint<Variant>("T"), ArgOp); \
\ \
REGISTER_KERNEL_BUILDER(Name(kRetOp) \ REGISTER_KERNEL_BUILDER( \
.Device(DEVICE) \ Name(kRetOp).Device(DEVICE).TypeConstraint("T", TYPES), RetvalOp); \
.TypeConstraint("T", TYPES) \
.HostMemory("input"), \
RetvalOp); \
REGISTER_KERNEL_BUILDER(Name(kRetOp) \ REGISTER_KERNEL_BUILDER(Name(kRetOp) \
.Device(DEVICE) \ .Device(DEVICE) \
.TypeConstraint<ResourceHandle>("T") \ .TypeConstraint<ResourceHandle>("T") \