Home
last modified time | relevance | path

Searched refs:suffix_dim_size (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/
Dsplit_op.cc122 IndexType suffix_dim_size = 1; in SetDims() local
124 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i)); in SetDims()
126 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size); in SetDims()
139 Eigen::DenseIndex suffix_dim_size, in operator ()() argument
160 split_dim_output_size, suffix_dim_size, &sizes, in operator ()()
167 if (prefix_dim_size * split_dim_output_size * suffix_dim_size > 0) { in operator ()()
230 Eigen::DenseIndex suffix_dim_size; in Compute() local
232 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) = in Compute()
239 input.shaped<T, 2>({split_dim_size, suffix_dim_size}); in Compute()
241 return Eigen::DSizes<Eigen::DenseIndex, 2>{split_size, suffix_dim_size}; in Compute()
[all …]
Dsplit_lib_gpu.cu.cc79 int32 split_dim_size, int32 suffix_dim_size, in SplitOpKernel() argument
88 int32 size = prefix_dim_size * split_dim_size * suffix_dim_size; in SplitOpKernel()
93 int32 i = offset / (split_dim_size * suffix_dim_size); in SplitOpKernel()
94 int32 j = (offset % (split_dim_size * suffix_dim_size)) / suffix_dim_size; in SplitOpKernel()
95 int32 k = offset % suffix_dim_size; in SplitOpKernel()
104 int32 output_offset = i * piece_size * suffix_dim_size + in SplitOpKernel()
105 (j % piece_size) * suffix_dim_size + k; in SplitOpKernel()
174 int32 suffix_dim_size, in SplitVOpKernel_fixed() argument
182 int32 size = prefix_dim_size * suffix_dim_size; in SplitVOpKernel_fixed()
183 int32 piece_size = suffix_dim_size / num_split; in SplitVOpKernel_fixed()
[all …]
Dsplit_v_op.cc176 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 …]
Done_hot_op.cc106 TI suffix_dim_size = indices_shape.num_elements() / prefix_dim_size; in Compute() local
110 indices.shaped<TI, 2>({prefix_dim_size, suffix_dim_size}); in Compute()
114 output->shaped<T, 3>({prefix_dim_size, depth_v, suffix_dim_size}); in Compute()
Dsplit_lib_gpu.h35 int32 split_dim_size, int32 suffix_dim_size,
Dbatch_kernels.cc145 int64 suffix_dim_size = 1; in SplitCPU() local
147 suffix_dim_size *= input.shape().dim_size(i); in SplitCPU()
150 input.shaped<T, 2>({input.shape().dim_size(0), suffix_dim_size}); in SplitCPU()
159 auto output_shaped = output.shaped<T, 2>({size, suffix_dim_size}); in SplitCPU()
162 Eigen::DSizes<Eigen::DenseIndex, 2> slice_sizes{size, suffix_dim_size}; in SplitCPU()
/external/tensorflow/tensorflow/lite/kernels/
Done_hot.cc70 const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size; in OneHotComputeImpl() local
86 for (int k = 0; k < suffix_dim_size; ++k, ++output) { in OneHotComputeImpl()
87 *output = static_cast<int>(indices[i * suffix_dim_size + k]) == j in OneHotComputeImpl()