From ef9eb1cf64ca9e09d7c3c9fa67eb6a9ea5985658 Mon Sep 17 00:00:00 2001
From: Sung Jin Hwang <sjhwang@google.com>
Date: Thu, 31 Oct 2019 15:53:55 -0700
Subject: [PATCH] Add int32 type support for BiasAdd. CPU support already
exists, so adding GPU support only.
PiperOrigin-RevId: 277818665
Change-Id: I174cdab534dc286e7ddd221e95bdef73e372097c
---
tensorflow/core/kernels/bias_op.cc | 1 +
tensorflow/core/kernels/bias_op_gpu.cu.cc | 3 +++
tensorflow/python/kernel_tests/bias_op_base.py | 2 +-
3 files changed, 5 insertions(+), 1 deletion(-)
diff --git a/tensorflow/core/kernels/bias_op.cc b/tensorflow/core/kernels/bias_op.cc
index 00b920c5f45..bf001cceae7 100644
--- a/tensorflow/core/kernels/bias_op.cc
+++ b/tensorflow/core/kernels/bias_op.cc
@@ -379,6 +379,7 @@ class BiasOp<GPUDevice, T> : public BinaryOp<T> {
BiasOp<GPUDevice, type>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL);
+REGISTER_GPU_KERNEL(int32);
#undef REGISTER_GPU_KERNEL
struct BiasGradAutotuneGroup {
diff --git a/tensorflow/core/kernels/bias_op_gpu.cu.cc b/tensorflow/core/kernels/bias_op_gpu.cu.cc
index bf54aff5a0d..54f9505b459 100644
--- a/tensorflow/core/kernels/bias_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/bias_op_gpu.cu.cc
@@ -289,6 +289,9 @@ void BiasGradGPU<T>::DoColReduction(OpKernelContext* context, T* output,
TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPECS);
+// No BiasGrad kernel for int32.
+template struct BiasGPU<int32>;
+
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
diff --git a/tensorflow/python/kernel_tests/bias_op_base.py b/tensorflow/python/kernel_tests/bias_op_base.py
index 9b5d4131f6b..cb7e771febf 100644
--- a/tensorflow/python/kernel_tests/bias_op_base.py
+++ b/tensorflow/python/kernel_tests/bias_op_base.py
@@ -83,7 +83,7 @@ class BiasAddTestBase(test.TestCase):
def _testAll(self, np_inputs, np_bias):
self._testBias(np_inputs, np_bias, use_gpu=False)
self._testBiasNCHW(np_inputs, np_bias, use_gpu=False)
- if np_inputs.dtype in [np.float16, np.float32, np.float64]:
+ if np_inputs.dtype in [np.float16, np.float32, np.float64, np.int32]:
self._testBias(np_inputs, np_bias, use_gpu=True)
self._testBiasNCHW(np_inputs, np_bias, use_gpu=True)