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 | 103 split_tensor = xla_sharding.split(tensor, 2, 3)
104 self.assertIsInstance(split_tensor, ops.Tensor)
105 split_sharding = xla_sharding.get_tensor_sharding(split_tensor)
109 return split_tensor
120 split_tensor = xla_sharding.split(tensor, 0, 8)
121 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 | 590 std::vector<Tensor> split_tensor; in SplitOutputTensors() local
592 output_tensor, task_sizes_plus_optional_padding, &split_tensor); in SplitOutputTensors()
598 DCHECK_EQ(split_tensor.size(), task_sizes_plus_optional_padding.size()); in SplitOutputTensors()
599 if (split_tensor.size() != task_sizes_plus_optional_padding.size()) { in SplitOutputTensors()
602 split_tensor.size(), " splits; expected ", in SplitOutputTensors()
611 tensor_vector[i] = std::move(split_tensor[j]); in SplitOutputTensors()
613 task.context->set_output(i, split_tensor[j]); in SplitOutputTensors()
|