| /external/tensorflow/tensorflow/core/kernels/ |
| D | gather_nd_op_gpu.cu.cc | 29 template <typename T, typename Index, int IXDIM>
32 T* __restrict__ out, const Eigen::array<int64, IXDIM> batch_strides, in GatherSliceOpKernel()
33 const Eigen::array<int64, IXDIM> batch_indices, const int64 indices_size, in GatherSliceOpKernel()
39 const auto indices_i = indices + IXDIM * loc; in GatherSliceOpKernel()
43 for (int j = 0; j < IXDIM; ++j) { in GatherSliceOpKernel()
67 template <typename T, typename Index, int IXDIM>
68 struct GatherNdSlice<GPUDevice, T, Index, IXDIM> {
71 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams, in operator ()()
77 Eigen::array<int64, IXDIM> batch_strides; in operator ()()
78 Eigen::array<int64, IXDIM> batch_indices; in operator ()()
[all …]
|
| D | gather_nd_op_cpu_impl.h | 41 template <typename T, typename Index, int IXDIM>
46 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams, in GatherNdSliceGenerator()
55 const Index loc, Eigen::array<Eigen::DenseIndex, IXDIM + 1>* ix) const { in GenerateIndices()
56 (*ix)[IXDIM] = 0; in GenerateIndices()
58 for (int i = 0; i < IXDIM; ++i) { in GenerateIndices()
69 Eigen::array<Eigen::DenseIndex, IXDIM + 1> ix; in operator()
87 const typename TTypes<T, IXDIM + 1>::ConstTensor Tparams_;
96 template <typename T, typename Index, int IXDIM>
97 struct GatherNdSlice<CPUDevice, T, Index, IXDIM> {
100 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams,
[all …]
|
| D | scatter_nd_op_gpu.cu.cc | 97 template <typename T, typename Index, scatter_nd_op::UpdateOp op, int IXDIM>
100 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix, in ScatterNdOpKernel()
101 const Eigen::array<int64, IXDIM> batch_strides, const int64 num_indices, in ScatterNdOpKernel()
109 for (int dim = 0; dim < IXDIM; ++dim) { in ScatterNdOpKernel()
110 int offset = (IXDIM * index + dim); in ScatterNdOpKernel()
127 template <typename T, typename Index, scatter_nd_op::UpdateOp op, int IXDIM>
128 struct ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM> {
131 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix, in operator ()()
144 Eigen::array<int64, IXDIM> batch_strides; in operator ()()
145 for (int dim = IXDIM - 1; dim >= 0; --dim) { in operator ()()
[all …]
|
| D | gather_nd_op.h | 36 template <typename Device, typename T, typename Index, int IXDIM>
43 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams,
133 #define PARAMS_CASE(IXDIM) \ in DoGatherNd() argument
134 case IXDIM: { \ in DoGatherNd()
135 functor::GatherNdSlice<Device, T, Index, IXDIM> func; \ in DoGatherNd()
136 auto params_flat = params.flat_outer_dims<T, IXDIM + 1>(); \ in DoGatherNd()
|
| D | scatter_nd_op_cpu_impl.h | 106 template <typename T, typename Index, scatter_nd_op::UpdateOp OP, int IXDIM>
107 struct ScatterNdFunctor<CPUDevice, T, Index, OP, IXDIM> {
110 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix,
121 Index batch_strides[IXDIM];
122 for (int dim = IXDIM - 1; dim >= 0; --dim) {
123 if (dim == IXDIM - 1) {
134 for (int dim = 0; dim < IXDIM; ++dim) {
|
| D | scatter_nd_op.h | 48 scatter_nd_op::UpdateOp op, int IXDIM>
53 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix,
|
| D | scatter_nd_op.cc | 942 #define PARAMS_CASE(IXDIM) \ in DoScatterNd() argument
943 case IXDIM: { \ in DoScatterNd()
944 typename Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix; \ in DoScatterNd()
945 for (int i = 0; i < IXDIM; ++i) { \ in DoScatterNd()
948 functor::ScatterNdFunctor<Device, T, Index, Op, IXDIM> functor; \ in DoScatterNd()
986 #define DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, IXDIM) \ argument
988 Index ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM>::operator()( \
990 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix, \
995 extern template struct ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM>;
|