| /external/tensorflow/tensorflow/core/kernels/ |
| split_op.cc | 112 int32 prefix_dim_size = 1;
114 prefix_dim_size *= input_shape.dim_size(i);
126 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
137 Eigen::DenseIndex prefix_dim_size,
159 auto range_output_func = [&indices, context, &output_shape, prefix_dim_size,
167 if (prefix_dim_size * split_dim_output_size * suffix_dim_size > 0) {
228 Eigen::DenseIndex prefix_dim_size;
232 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
237 if (prefix_dim_size == 1) {
247 context, input_reshaped, input_shape, split_dim, prefix_dim_size, member in class:SplitOpCPU
262 context, input_reshaped, input_shape, split_dim, prefix_dim_size, member in class:__anon45533
[all...] |
| split_v_op.cc | 166 int32 prefix_dim_size = 1; local
168 prefix_dim_size *= input_shape.dim_size(i);
180 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
192 Eigen::DenseIndex prefix_dim_size,
283 Eigen::DenseIndex prefix_dim_size;
287 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
299 if (prefix_dim_size == 1) {
310 prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec,
314 {prefix_dim_size, split_dim_size, suffix_dim_size});
316 return Eigen::DSizes<Eigen::DenseIndex, 3>{prefix_dim_size, split_size
325 prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec, member in class:tensorflow::SplitVOpCPU
[all...] |
| one_hot_op.cc | 99 // prefix_dim_size == # of elements before the axis
102 int64 prefix_dim_size = 1; variable
104 prefix_dim_size *= indices_shape.dim_size(i);
106 TI suffix_dim_size = indices_shape.num_elements() / prefix_dim_size;
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});
|
| split_lib_gpu.h | 34 void Run(const Eigen::GpuDevice& d, const T* input, int32 prefix_dim_size,
|
| split_lib_gpu.cu.cc | 78 __global__ void SplitOpKernel(const T* input, int32 prefix_dim_size,
88 int32 size = prefix_dim_size * split_dim_size * suffix_dim_size;
100 // [prefix_dim_size][piece_size][suffix_dim_size].
173 __global__ void SplitVOpKernel_fixed(const T* input, int32 prefix_dim_size,
182 int32 size = prefix_dim_size * suffix_dim_size;
199 int32 prefix_dim_size, int32 split_dim_size,
203 prefix_dim_size * split_dim_size * suffix_dim_size, d);
207 prefix_dim_size, split_dim_size, suffix_dim_size,
|
| /external/tensorflow/tensorflow/lite/kernels/ |
| one_hot.cc | 63 // prefix_dim_size == # of elements before the axis
66 int prefix_dim_size = 1; local
68 prefix_dim_size *= op_context.indices->dims->data[i];
70 const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size;
77 // prefix_dim_size x suffix_dim_size
79 // prefix_dim_size x depth x suffix_dim_size
84 for (int i = 0; i < prefix_dim_size; ++i) {
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| one_hot_op_test.py | 370 prefix_dim_size = 65536
372 x = [i % depth for i in range(prefix_dim_size)]
375 truth = np.zeros((prefix_dim_size, depth), np.float32)
376 for i in range(prefix_dim_size):
|
