Searched refs:split_tensor (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | split_v_op.cc | 58 const Tensor& split_tensor = context->input(1); in ComputeEasyCases() local
71 split_tensor.dims() == 1 && split_tensor.NumElements() == num_split, in ComputeEasyCases()
74 split_tensor.dims(), " -D and ", in ComputeEasyCases()
75 split_tensor.NumElements(), " elements")); in ComputeEasyCases()
77 auto split_sizes_d = split_tensor.vec<Tlen>(); in ComputeEasyCases()
|
| D | batch_kernels.cc | 447 std::vector<Tensor> split_tensor; in SplitOutputTensors() local
449 output_tensor, task_sizes_plus_optional_padding, &split_tensor); in SplitOutputTensors()
455 DCHECK_EQ(split_tensor.size(), task_sizes_plus_optional_padding.size()); in SplitOutputTensors()
456 if (split_tensor.size() != task_sizes_plus_optional_padding.size()) { in SplitOutputTensors()
459 split_tensor.size(), " splits; expected ", in SplitOutputTensors()
465 task.context->set_output(i, split_tensor.at(j)); in SplitOutputTensors()
|