| /external/tensorflow/tensorflow/core/kernels/ |
| gather_nd_op_gpu.cu.cc | 29 template <typename T, typename Index, int IXDIM>
32 const Eigen::array<int64, IXDIM> batch_strides,
33 const Eigen::array<int64, IXDIM> batch_indices, const int64 indices_size,
39 const auto indices_i = indices + IXDIM * loc;
43 for (int j = 0; j < IXDIM; ++j) {
67 template <typename T, typename Index, int IXDIM>
68 struct GatherNdSlice<GPUDevice, T, Index, IXDIM> {
71 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams,
77 Eigen::array<int64, IXDIM> batch_strides;
78 Eigen::array<int64, IXDIM> batch_indices
[all...] |
| scatter_nd_op_gpu.cu.cc | 81 template <typename T, typename Index, scatter_nd_op::UpdateOp op, int IXDIM>
84 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix,
85 const Eigen::array<int64, IXDIM> batch_strides, const int64 num_indices,
93 for (int dim = 0; dim < IXDIM; ++dim) {
94 int offset = (IXDIM * index + dim);
111 template <typename T, typename Index, scatter_nd_op::UpdateOp op, int IXDIM>
112 struct ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM> {
115 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix,
127 // Index batch_strides[IXDIM];
128 Eigen::array<int64, IXDIM> batch_strides
[all...] |
| gather_nd_op_cpu_impl.h | 41 template <typename T, typename Index, int IXDIM>
46 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams,
55 const Index loc, Eigen::array<Eigen::DenseIndex, IXDIM + 1>* ix) const {
56 (*ix)[IXDIM] = 0;
58 for (int i = 0; i < IXDIM; ++i) {
69 Eigen::array<Eigen::DenseIndex, IXDIM + 1> ix;
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...] |
| scatter_nd_op_cpu_impl.h | 90 template <typename T, typename Index, scatter_nd_op::UpdateOp OP, int IXDIM>
91 struct ScatterNdFunctor<CPUDevice, T, Index, OP, IXDIM> {
94 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix,
105 Index batch_strides[IXDIM];
106 for (int dim = IXDIM - 1; dim >= 0; --dim) {
107 if (dim == IXDIM - 1) {
118 for (int dim = 0; dim < IXDIM; ++dim) {
172 template <typename T, typename Index, scatter_nd_op::UpdateOp OP, int IXDIM>
173 struct ScatterNdFunctor<SYCLDevice, T, Index, OP, IXDIM> {
176 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix
[all...] |
| gather_nd_op.h | 32 template <typename Device, typename T, typename Index, int IXDIM>
39 typename TTypes<T, IXDIM + 1>::ConstTensor Tparams,
|
| scatter_nd_op.cc | 472 #define PARAMS_CASE(IXDIM) \
473 case IXDIM: { \
474 typename Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix; \
475 for (int i = 0; i < IXDIM; ++i) { \
478 functor::ScatterNdFunctor<Device, T, Index, Op, IXDIM> functor; \
514 #define DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, IXDIM) \
516 Index ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM>::operator()( \
518 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix, \
523 extern template struct ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM>;
|
| scatter_nd_op.h | 48 scatter_nd_op::UpdateOp op, int IXDIM>
53 const Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix,
|
| gather_nd_op.cc | 166 #define PARAMS_CASE(IXDIM) \
167 case IXDIM: { \
168 functor::GatherNdSlice<Device, T, Index, IXDIM> func; \
169 auto params_flat = params.flat_outer_dims<T, IXDIM + 1>(); \
|
| /external/tensorflow/tensorflow/compiler/tests/ |
| scatter_nd_op_test.py | 50 ixdim = indices.shape[-1]
51 num_updates = indices.size // ixdim
54 for i in range(ixdim, total_nd):
58 output_flat = _FlatOuterDims(ref, ixdim + 1)
84 ixdim = indices_shape[-1]
87 for i in range(ixdim):
105 for i in range(ixdim, len(ref_shape)):
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| scatter_nd_ops_test.py | 55 ixdim = indices.shape[-1]
56 num_updates = indices.size // ixdim
59 for i in range(ixdim, total_nd):
63 output_flat = _FlatOuterDims(ref, ixdim + 1)
105 ixdim = indices_shape[-1]
108 for i in range(ixdim):
127 for i in range(ixdim, len(ref_shape)):
|
| /external/tensorflow/tensorflow/core/ops/ |
| array_ops.cc | 2632 const DimensionHandle ixdim = c->Dim(indices_shape, -1); local
[all...] |
