| /external/tensorflow/tensorflow/core/kernels/ |
| split_lib_gpu.cu.cc | 79 int32 split_dim_size, int32 suffix_dim_size,
88 int32 size = prefix_dim_size * split_dim_size * suffix_dim_size;
93 int32 i = offset / (split_dim_size * suffix_dim_size);
94 int32 j = (offset % (split_dim_size * suffix_dim_size)) / suffix_dim_size;
95 int32 k = offset % suffix_dim_size;
100 // [prefix_dim_size][piece_size][suffix_dim_size].
104 int32 output_offset = i * piece_size * suffix_dim_size +
105 (j % piece_size) * suffix_dim_size + k;
174 int32 suffix_dim_size,
[all...] |
| split_op.cc | 117 // Caller must ensure that dim_size and suffix_dim_size are <
122 IndexType suffix_dim_size = 1;
124 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i));
126 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
139 Eigen::DenseIndex suffix_dim_size,
160 split_dim_output_size, suffix_dim_size, &sizes,
167 if (prefix_dim_size * split_dim_output_size * suffix_dim_size > 0) {
230 Eigen::DenseIndex suffix_dim_size;
232 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
239 input.shaped<T, 2>({split_dim_size, suffix_dim_size});
248 split_dim_size, suffix_dim_size, make_sizes, reshape_result, member in class:SplitOpCPU
263 split_dim_size, suffix_dim_size, make_sizes, reshape_result, member in class:__anon45533
[all...] |
| split_v_op.cc | 171 // Caller must ensure that dim_size and suffix_dim_size are <
176 IndexType suffix_dim_size = 1; local
178 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i));
180 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
194 Eigen::DenseIndex suffix_dim_size,
285 Eigen::DenseIndex suffix_dim_size;
287 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
301 input.shaped<T, 2>({split_dim_size, suffix_dim_size});
303 return Eigen::DSizes<Eigen::DenseIndex, 2>{split_size, suffix_dim_size};
306 return result->shaped<T, 2>({split_size, suffix_dim_size});
325 prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec, member in class:tensorflow::SplitVOpCPU
[all...] |
| split_lib_gpu.h | 35 int32 split_dim_size, int32 suffix_dim_size,
|
| one_hot_op.cc | 101 // suffix_dim_size == # of elements after the axis
106 TI suffix_dim_size = indices_shape.num_elements() / prefix_dim_size; variable
108 // Split indices into matrix of size prefix_dim_size x suffix_dim_size
110 indices.shaped<TI, 2>({prefix_dim_size, suffix_dim_size});
112 // prefix_dim_size x depth x suffix_dim_size.
114 output->shaped<T, 3>({prefix_dim_size, depth_v, suffix_dim_size});
|
| batch_kernels.cc | 145 int64 suffix_dim_size = 1; local
147 suffix_dim_size *= input.shape().dim_size(i);
150 input.shaped<T, 2>({input.shape().dim_size(0), suffix_dim_size});
159 auto output_shaped = output.shaped<T, 2>({size, suffix_dim_size});
162 Eigen::DSizes<Eigen::DenseIndex, 2> slice_sizes{size, suffix_dim_size};
[all...] |
| /external/tensorflow/tensorflow/lite/kernels/ |
| one_hot.cc | 65 // suffix_dim_size == # of elements after the axis
70 const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size; local
77 // prefix_dim_size x suffix_dim_size
79 // prefix_dim_size x depth x suffix_dim_size
86 for (int k = 0; k < suffix_dim_size; ++k, ++output) {
87 *output = static_cast<int>(indices[i * suffix_dim_size + k]) == j
|
