diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2017-02-02 18:51:45 -0800 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2017-02-02 19:06:31 -0800 |
| commit | ef489fb30b374271da1b4a5ef549a9965737dcbb (patch) | |
| tree | 1fb07437ff7a5b2daab11d87bd488f110fb15e14 /tensorflow/core/kernels/unique_op.cc | |
| parent | ec1dec819c7ef5540a45ac75e2278d4f7ae0e9ad (diff) | |
Fix for op placements when deduplicating sparse gradients.
Adds a fake GPU kernel for unique, which allows the hard placement of other ops on the GPU.
Change: 146432826
Diffstat (limited to 'tensorflow/core/kernels/unique_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/unique_op.cc | 12 |
1 files changed, 12 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/unique_op.cc b/tensorflow/core/kernels/unique_op.cc index 4a2a60bdfb..9f44ee5b19 100644 --- a/tensorflow/core/kernels/unique_op.cc +++ b/tensorflow/core/kernels/unique_op.cc @@ -94,4 +94,16 @@ class UniqueOp : public OpKernel { TF_CALL_REAL_NUMBER_TYPES(REGISTER_UNIQUE); REGISTER_UNIQUE(string) #undef REGISTER_UNIQUE + +// A fake int32 GPU kernel so that the use of Unique in optimizers (to +// de-duplicate sparse gradient indices) does not conflict with gradients being +// located on a GPU. +REGISTER_KERNEL_BUILDER(Name("Unique") + .Device(DEVICE_GPU) + .TypeConstraint<int32>("T") + .TypeConstraint<int32>("out_idx") + .HostMemory("x") + .HostMemory("y") + .HostMemory("idx"), + UniqueOp<int32>); } // namespace tensorflow |