Searched refs:suffix_dim_size (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | split_op.cc | 119 IndexType suffix_dim_size = 1; in SetDims() local
121 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i)); in SetDims()
123 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size); in SetDims()
136 Eigen::DenseIndex suffix_dim_size, in operator ()() argument
157 split_dim_output_size, suffix_dim_size, &sizes, in operator ()()
164 if (prefix_dim_size * split_dim_output_size * suffix_dim_size > 0) { in operator ()()
227 Eigen::DenseIndex suffix_dim_size; in Compute() local
229 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) = in Compute()
236 input.shaped<T, 2>({split_dim_size, suffix_dim_size}); in Compute()
238 return Eigen::DSizes<Eigen::DenseIndex, 2>{split_size, suffix_dim_size}; in Compute()
[all …]
|
| D | split_lib_gpu.cu.cc | 75 int32 suffix_dim_size, in SplitOpKernel() argument
84 int32 size = prefix_dim_size * split_dim_size * suffix_dim_size; in SplitOpKernel()
89 int32 i = offset / (split_dim_size * suffix_dim_size); in SplitOpKernel()
90 int32 j = (offset % (split_dim_size * suffix_dim_size)) / suffix_dim_size; in SplitOpKernel()
91 int32 k = offset % suffix_dim_size; in SplitOpKernel()
100 int32 output_offset = i * piece_size * suffix_dim_size + in SplitOpKernel()
101 (j % piece_size) * suffix_dim_size + k; in SplitOpKernel()
171 int32 suffix_dim_size, in SplitVOpKernel_fixed() argument
179 int32 size = prefix_dim_size * suffix_dim_size; in SplitVOpKernel_fixed()
180 int32 piece_size = suffix_dim_size / num_split; in SplitVOpKernel_fixed()
[all …]
|
| D | split_v_op.cc | 176 IndexType suffix_dim_size = 1; in SetDims() local
178 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i)); in SetDims()
180 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size); in SetDims()
194 Eigen::DenseIndex suffix_dim_size, in operator ()() argument
285 Eigen::DenseIndex suffix_dim_size; in Compute() local
287 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) = in Compute()
301 input.shaped<T, 2>({split_dim_size, suffix_dim_size}); in Compute()
303 return Eigen::DSizes<Eigen::DenseIndex, 2>{split_size, suffix_dim_size}; in Compute()
306 return result->shaped<T, 2>({split_size, suffix_dim_size}); in Compute()
310 prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec, in Compute()
[all …]
|
| D | one_hot_op.cc | 107 int64 suffix_dim_size = indices_shape.num_elements() / prefix_dim_size; in Compute() local
111 indices.shaped<TI, 2>({prefix_dim_size, suffix_dim_size}); in Compute()
115 output->shaped<T, 3>({prefix_dim_size, depth_v, suffix_dim_size}); in Compute()
|
| D | split_lib_gpu.h | 35 int32 split_dim_size, int32 suffix_dim_size,
|
| /external/tensorflow/tensorflow/core/kernels/batching_util/ |
| D | concat_split_util.h | 160 int64 suffix_dim_size = 1; in SplitCPU() local
162 suffix_dim_size *= input.shape().dim_size(i); in SplitCPU()
165 input.shaped<T, 2>({input.shape().dim_size(0), suffix_dim_size}); in SplitCPU()
176 auto output_shaped = output.shaped<T, 2>({size, suffix_dim_size}); in SplitCPU()
182 static_cast<Eigen::DenseIndex>(suffix_dim_size)}; in SplitCPU()
|
| /external/tensorflow/tensorflow/lite/kernels/ |
| D | one_hot.cc | 72 const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size; in OneHotComputeImpl() local
88 for (int k = 0; k < suffix_dim_size; ++k, ++output) { in OneHotComputeImpl()
89 *output = static_cast<int>(indices[i * suffix_dim_size + k]) == j in OneHotComputeImpl()
|