Searched refs:gather_dim_size (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | gather_functor_gpu.cu.h | 35 T* __restrict__ out, int64 gather_dim_size, in GatherOpKernel() argument
59 if (!FastBoundsCheck(gather_i, gather_dim_size)) { in GatherOpKernel()
68 (batch_i * gather_dim_size + gather_i) * slice_size + slice_i; in GatherOpKernel()
91 const int64 gather_dim_size = params.dimension(1);
102 out.data(), gather_dim_size, indices_size, slice_size, out_size));
110 out.data(), gather_dim_size, indices_size, slice_size, out_size));
|
| D | gather_functor_batched_gpu.cu.h | 36 int64 gather_dim_size, int64 indices_size, in GatherOpKernel() argument
70 if (!FastBoundsCheck(gather_i, gather_dim_size)) { in GatherOpKernel()
78 (batch_i * outer_size + outer_i) * gather_dim_size + gather_i in GatherOpKernel()
104 const int64 gather_dim_size = params.dimension(2);
119 outer_size, gather_dim_size, indices_size, slice_size, out_size));
|
| D | gather_op.cc | 125 int64 gather_dim_size = params.dim_size(axis); in Compute() local
128 c, gather_dim_size <= std::numeric_limits<Index>::max(), in Compute()
131 " indexing: ", gather_dim_size, " > ", in Compute()
166 {batch_size, outer_size, gather_dim_size, inner_size}); in Compute()
174 params.shaped<T, 3>({outer_size, gather_dim_size, inner_size}); in Compute()
184 indices_flat(bad_i), " is not in [0, ", gather_dim_size, ")")); in Compute()
|
| D | dynamic_partition_op_gpu.cu.cc | 110 T* out, int64 gather_dim_size, int64 indices_size, in CallGatherKernel() argument
115 d.stream(), params, indices, out, gather_dim_size, in CallGatherKernel()
|
| D | resource_variable_ops.cc | 685 int64 gather_dim_size = 1; in Compute() local
687 gather_dim_size *= params.dim_size(idx); in Compute()
693 auto params_flat = params.shaped<T, 3>({1, gather_dim_size, inner_size}); in Compute()
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | array_grad.py | 607 gather_dim_size): argument
619 gather_dim_size *= batch_size
626 params_grad = math_ops.unsorted_segment_sum(values, indices, gather_dim_size)
|