| /external/tensorflow/tensorflow/core/kernels/ |
| D | split_op.cc | 116 IndexType split_dim_size = in SetDims() local
123 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size); in SetDims()
135 Eigen::DenseIndex split_dim_size, in operator ()() argument
226 Eigen::DenseIndex split_dim_size; in Compute() local
229 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) = in Compute()
232 const int64 split_dim_output_size = split_dim_size / num_split; in Compute()
236 input.shaped<T, 2>({split_dim_size, suffix_dim_size}); in Compute()
245 split_dim_size, suffix_dim_size, make_sizes, reshape_result, in Compute()
249 {prefix_dim_size, split_dim_size, suffix_dim_size}); in Compute()
260 split_dim_size, suffix_dim_size, make_sizes, reshape_result, in Compute()
[all …]
|
| D | split_lib_gpu.cu.cc | 74 int32 prefix_dim_size, int32 split_dim_size, in SplitOpKernel() argument
82 eigen_assert(split_dim_size % num_split == 0); in SplitOpKernel()
84 int32 size = prefix_dim_size * split_dim_size * suffix_dim_size; in SplitOpKernel()
85 int32 piece_size = split_dim_size / num_split; 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()
196 int32 prefix_dim_size, int32 split_dim_size, in Run() argument
200 GetGpuLaunchConfig(prefix_dim_size * split_dim_size * suffix_dim_size, d); in Run()
204 prefix_dim_size, split_dim_size, suffix_dim_size, in Run()
|
| D | split_v_op.cc | 173 IndexType split_dim_size = in SetDims() local
180 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size); in SetDims()
193 Eigen::DenseIndex split_dim_size, in operator ()() argument
284 Eigen::DenseIndex split_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()
310 prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec, in Compute()
314 {prefix_dim_size, split_dim_size, suffix_dim_size}); in Compute()
325 prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec, in Compute()
360 int32 split_dim_size; in Compute() local
[all …]
|
| D | split_lib_gpu.h | 35 int32 split_dim_size, int32 suffix_dim_size,
|
| /external/tensorflow/tensorflow/core/util/sparse/ |
| D | sparse_tensor.h | 514 const int split_dim_size = input_tensor.shape()[split_dim]; in Split() local
515 const int split_size = split_dim_size / num_split; in Split()
517 if (!(num_split > 0 && num_split <= split_dim_size)) { in Split()
519 split_dim_size, "]"); in Split()
526 const int residual = split_dim_size % num_split; in Split()
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | array_ops.cc | 589 DimensionHandle split_dim_size; in __anondb9326b20a02() local
592 true /* evenly_divisible */, &split_dim_size), in __anondb9326b20a02()
595 c->ReplaceDim(input, split_dim, split_dim_size, &out)); in __anondb9326b20a02()
673 auto split_dim_size = c->Value(c->Dim(input, split_dim)); in __anondb9326b20b02() local
680 if (data[i] == -1 && c->ValueKnown(split_dim_size)) { in __anondb9326b20b02()
681 size = split_dim_size - total_size; in __anondb9326b20b02()
687 if (c->ValueKnown(split_dim_size)) { in __anondb9326b20b02()
688 if (has_neg_one ? total_size > split_dim_size in __anondb9326b20b02()
689 : total_size != split_dim_size) { in __anondb9326b20b02()
691 "can't split axis of size ", split_dim_size, in __anondb9326b20b02()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/IR/ |
| D | hlo_ops.cc | 682 auto split_dim_size = type.getDimSize(op.split_dimension()); in Verify() local
684 if (split_dim_size % split_count != 0) { in Verify()
685 return op.emitError() << "split dimension has size " << split_dim_size in Verify()
|