1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
|
#include "tensorflow/core/kernels/cwise_ops_common.h"
namespace tensorflow {
REGISTER4(UnaryOp, CPU, "Sign", functor::sign, float, double, int32, int64);
#if GOOGLE_CUDA
REGISTER3(UnaryOp, GPU, "Sign", functor::sign, float, double, int64);
#endif
// A special GPU kernel for int32.
// TODO(b/25387198): Also enable int32 in device memory. This kernel
// registration requires all int32 inputs and outputs to be in host memory.
REGISTER_KERNEL_BUILDER(Name("Sign")
.Device(DEVICE_GPU)
.HostMemory("x")
.HostMemory("y")
.TypeConstraint<int32>("T"),
UnaryOp<CPUDevice, functor::sign<int32>>);
} // namespace tensorflow
|
