Searched refs:gather_dim_size (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | gather_functor_gpu.cu.h | 34 int64 gather_dim_size, int64 indices_size, in GatherOpKernel() argument
57 if (!FastBoundsCheck(gather_i, gather_dim_size)) { in GatherOpKernel()
66 (batch_i * gather_dim_size + gather_i) * slice_size + slice_i; in GatherOpKernel()
89 const int64 gather_dim_size = params.dimension(1);
98 out.data(), gather_dim_size, indices_size, slice_size, out_size));
103 out.data(), gather_dim_size, indices_size, slice_size, out_size));
|
| D | gather_op.cc | 78 const int64 gather_dim_size = params.dim_size(axis); in Compute() local
81 c, gather_dim_size <= std::numeric_limits<Index>::max(), in Compute()
84 " indexing: ", gather_dim_size, " > ", in Compute()
106 params.shaped<T, 3>({outer_size, gather_dim_size, inner_size}); in Compute()
117 indices_flat(bad_i), " is not in [0, ", gather_dim_size, ")")); in Compute()
|
| D | dynamic_partition_op_gpu.cu.cc | 104 T* out, int64 gather_dim_size, int64 indices_size, in CallGatherKernel() argument
110 gather_dim_size, indices_size, slice_size, out_size)); in CallGatherKernel()
|
| D | resource_variable_ops.cc | 681 int64 gather_dim_size = 1; in Compute() local
683 gather_dim_size *= params.dim_size(idx); in Compute()
689 auto params_flat = params.shaped<T, 3>({1, gather_dim_size, inner_size}); in Compute()
|