diff options
| author |  Yong Tang <yong.tang.github@outlook.com> | 2017-09-26 01:12:37 +0000 |
|---|---|---|
| committer | Yong Tang <yong.tang.github@outlook.com> | 2017-09-26 01:19:21 +0000 |
| commit | 1398d4bec48618bc497b7be7cfd655d5d80c4efd (patch) | |
| tree | f112e359ddf7fab5f7badccab6859479a3c7f20c | |
| parent | 3537e663230ba2d694f6daebc2288f180c6db4ea (diff) | |
Add int64 support for `tf.reduce_max` on GPU
This fix tries to address the issue raised in 13293 where
`tf.reduce_max` on GPU does not have int64 support.
This fix fixes 13293.
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
| -rw-r--r-- | tensorflow/core/kernels/reduction_ops_gpu_int.cu.cc | 1 | ||||
| -rw-r--r-- | tensorflow/core/kernels/reduction_ops_max.cc | 1 |
2 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/reduction_ops_gpu_int.cu.cc b/tensorflow/core/kernels/reduction_ops_gpu_int.cu.cc index c2b4d05fe7..69296c7b65 100644 --- a/tensorflow/core/kernels/reduction_ops_gpu_int.cu.cc +++ b/tensorflow/core/kernels/reduction_ops_gpu_int.cu.cc @@ -59,6 +59,7 @@ typedef TTypes<float>::Tensor::Index Index; DEFINE_FOR_TYPE_AND_R(T, Eigen::internal::ProdReducer<T>) DEFINE_FOR_ALL_REDUCERS(int32); +DEFINE_FOR_ALL_REDUCERS(int64); #undef DEFINE_FOR_ALL_REDUCERS #undef DEFINE_FOR_TYPE_AND_R #undef DEFINE diff --git a/tensorflow/core/kernels/reduction_ops_max.cc b/tensorflow/core/kernels/reduction_ops_max.cc index d243e7c55f..4ca5c11a48 100644 --- a/tensorflow/core/kernels/reduction_ops_max.cc +++ b/tensorflow/core/kernels/reduction_ops_max.cc @@ -39,6 +39,7 @@ TF_CALL_REAL_NUMBER_TYPES(REGISTER_CPU_KERNELS); ReductionOp<GPUDevice, type, Eigen::internal::MaxReducer<type>>); REGISTER_GPU_KERNELS(float); REGISTER_GPU_KERNELS(double); +REGISTER_GPU_KERNELS(int64); // A special GPU kernel for int32. // TODO(b/25387198): Also enable int32 in device memory. This kernel |