Searched refs:split_tensor (Results 1 – 6 of 6) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| H A 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A 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()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/batching_util/ |
| H A 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()
|
| /aosp_15_r20/external/executorch/backends/vulkan/runtime/graph/ops/impl/ |
| H A D | Split.cpp | 132 void split_tensor(ComputeGraph& graph, const std::vector<ValueRef>& args) { in split_tensor() function
139 VK_REGISTER_OP(aten.split.Tensor, split_tensor);
|
| /aosp_15_r20/external/pytorch/test/ |
| H A D | test_flop_counter.py | 537 def split_tensor(x): function
543 q_tensors = split_tensor(q)
544 k_tensors = split_tensor(k)
545 v_tensors = split_tensor(v)
|
| /aosp_15_r20/external/pytorch/torch/nested/_internal/ |
| H A D | ops.py | 737 def split_tensor(func, *args, **kwargs): function
|