1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
|
#include "tensorflow/core/kernels/reduction_ops_common.h"
namespace tensorflow {
#define REGISTER_CPU_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Name("Sum").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
ReductionOp<CPUDevice, type, Eigen::internal::SumReducer<type>>);
TF_CALL_REAL_NUMBER_TYPES(REGISTER_CPU_KERNELS);
#undef REGISTER_CPU_KERNELS
// NOTE: We should have mean(complex64,int32), too. But that needs to
// change Eigen::internal::MeanReducer to cast int to complex<float>.
// We don't see immediate need of mean(complex64,int32) anyway.
REGISTER_KERNEL_BUILDER(
Name("Sum").Device(DEVICE_CPU).TypeConstraint<complex64>("T"),
ReductionOp<CPUDevice, complex64, Eigen::internal::SumReducer<complex64>>);
#if GOOGLE_CUDA
#define REGISTER_GPU_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Name("Sum") \
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T") \
.HostMemory("reduction_indices"), \
ReductionOp<GPUDevice, type, Eigen::internal::SumReducer<type>>);
REGISTER_GPU_KERNELS(float);
#undef REGISTER_GPU_KERNELS
REGISTER_KERNEL_BUILDER(
Name("Sum").Device(DEVICE_GPU).TypeConstraint<complex64>("T"),
ReductionOp<GPUDevice, complex64, Eigen::internal::SumReducer<complex64>>);
#endif
} // namespace tensorflow
|
