Searched refs:split_tensor (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding_test.py | 110 split_tensor = xla_sharding.split(tensor, 2, 3)
111 self.assertIsInstance(split_tensor, ops.Tensor)
112 split_sharding = xla_sharding.get_tensor_sharding(split_tensor)
116 return split_tensor
127 split_tensor = xla_sharding.split(tensor, 0, 8)
128 return split_tensor
|
| /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()
|
| /external/tensorflow/tensorflow/core/kernels/batching_util/ |
| D | batch_resource_base.cc | 626 std::vector<Tensor> split_tensor; in SplitOutputTensors() local
628 output_tensor, task_sizes_plus_optional_padding, &split_tensor); in SplitOutputTensors()
634 DCHECK_EQ(split_tensor.size(), task_sizes_plus_optional_padding.size()); in SplitOutputTensors()
635 if (split_tensor.size() != task_sizes_plus_optional_padding.size()) { in SplitOutputTensors()
638 split_tensor.size(), " splits; expected ", in SplitOutputTensors()
647 tensor_vector[i] = std::move(split_tensor[j]); in SplitOutputTensors()
649 task.context->set_output(i, split_tensor[j]); in SplitOutputTensors()
|