Home
last modified time | relevance | path

Searched refs:indices_size (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/
Dscatter_functor_gpu.cu.h76 Index indices_size) {
77 Index update_block = updates_size / indices_size;
96 Index indices_size,
98 Index update_block = synthesized_updates_size / indices_size;
127 const Index indices_size = indices.size();
133 indices.data(), first_dim_size, updates_size, indices_size));
148 const Index indices_size = indices.size();
149 const Index synthesized_updates_size = indices_size * params.dimension(1);
155 indices_size, synthesized_updates_size));
Dgather_functor_gpu.cu.h34 int64 gather_dim_size, int64 indices_size, in GatherOpKernel() argument
46 batch_i = batch_indices_i / indices_size; in GatherOpKernel()
48 indices_i = batch_indices_i - batch_i * indices_size; in GatherOpKernel()
90 const int64 indices_size = indices.size();
98 out.data(), gather_dim_size, indices_size, slice_size, out_size));
103 out.data(), gather_dim_size, indices_size, slice_size, out_size));
Dgather_functor.h44 const SliceIndex indices_size = static_cast<SliceIndex>(indices.dimension(0)); in HandleCopies() local
61 SliceIndex batch_idx = static_cast<SliceIndex>(start / indices_size); in HandleCopies()
62 SliceIndex indices_idx = static_cast<SliceIndex>(start % indices_size); in HandleCopies()
63 SliceIndex batch_idx_end = static_cast<SliceIndex>(end / indices_size); in HandleCopies()
64 SliceIndex indices_idx_end = static_cast<SliceIndex>(end % indices_size); in HandleCopies()
71 (i_next < indices_size)) { in HandleCopies()
93 out_base + (batch_idx * indices_size + indices_idx) * slice_elems, in HandleCopies()
108 batch_size * indices_size, slice_elems * sizeof(T), work); in HandleCopies()
Dgather_nd_op_gpu.cu.cc33 const Eigen::array<int64, IXDIM> batch_indices, const int64 indices_size, in GatherSliceOpKernel() argument
74 const int64 indices_size = Tindices.dimension(1); in operator ()() local
93 indices_size, s_size, out_size)); in operator ()()
Ddynamic_partition_op_gpu.cu.cc104 T* out, int64 gather_dim_size, int64 indices_size, in CallGatherKernel() argument
110 gather_dim_size, indices_size, slice_size, out_size)); in CallGatherKernel()
446 int32 indices_size = outs[p]->dim_size(0); in GatherSlices() local
451 indices_size, slice_size, out_size); in GatherSlices()
452 ind_base += indices_size; in GatherSlices()
/external/tensorflow/tensorflow/python/ops/
Darray_grad.py429 indices_size = array_ops.expand_dims(array_ops.size(indices), 0)
439 values_shape = array_ops.concat([indices_size, params_tail_shape], 0)
441 indices = array_ops.reshape(indices, indices_size)
453 values_shape = array_ops.concat([outer_shape, indices_size, inner_shape], 0)
455 indices = array_ops.reshape(indices, indices_size)
/external/tensorflow/tensorflow/python/kernel_tests/
Dgather_op_test.py416 indices_size = np.prod(indices_shape)
417 indices = np.reshape(np.arange(indices_size), indices_shape)
Dresource_variable_ops_test.py1214 indices_size = np.prod(indices_shape)
1215 indices = np.reshape(np.arange(indices_size, dtype=np.int32), indices_shape)