/external/tensorflow/tensorflow/core/kernels/ |
gather_functor_gpu.cu.h | 34 int64 gather_dim_size, int64 indices_size, 56 // Check gather_i is in [0, gather_dim_size). 57 if (!FastBoundsCheck(gather_i, gather_dim_size)) { 62 // params is a [batch_size, gather_dim_size, slice_size] tensor. Read 66 (batch_i * gather_dim_size + gather_i) * slice_size + slice_i; 89 const int64 gather_dim_size = params.dimension(1); local 98 out.data(), gather_dim_size, indices_size, slice_size, out_size)); 103 out.data(), gather_dim_size, indices_size, slice_size, out_size));
|
gather_op.cc | 78 const int64 gather_dim_size = params.dim_size(axis); variable 81 c, gather_dim_size <= std::numeric_limits<Index>::max(), 84 " indexing: ", gather_dim_size, " > ", 106 params.shaped<T, 3>({outer_size, gather_dim_size, inner_size}); 117 indices_flat(bad_i), " is not in [0, ", gather_dim_size, ")"));
|
dynamic_partition_op_gpu.cu.cc | 104 T* out, int64 gather_dim_size, int64 indices_size, 110 gather_dim_size, indices_size, slice_size, out_size));
|
resource_variable_ops.cc | 681 int64 gather_dim_size = 1; variable 683 gather_dim_size *= params.dim_size(idx); 689 auto params_flat = params.shaped<T, 3>({1, gather_dim_size, inner_size}); [all...] |