| /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/ |
| H A D | TensorShape.cu | 188 at::IntArrayRef split_sizes, in get_split_base_addrs() argument
195 split_base_addrs.reserve(split_sizes.size()); in get_split_base_addrs()
196 for (const auto& split_size : split_sizes) { in get_split_base_addrs()
215 at::IntArrayRef split_sizes, in get_split_chunk_sizes() argument
220 split_chunk_sizes.reserve(split_sizes.size()); in get_split_chunk_sizes()
221 for (const auto& split_size : split_sizes) { in get_split_chunk_sizes()
620 at::IntArrayRef split_sizes, in split_with_sizes_copy_out_cuda_contiguous_no_cast() argument
625 detail::get_split_base_addrs(self, split_sizes, dim); in split_with_sizes_copy_out_cuda_contiguous_no_cast()
629 detail::get_split_chunk_sizes(self, split_sizes, dim); in split_with_sizes_copy_out_cuda_contiguous_no_cast()
667 for (size_t split_idx = 0; split_idx < split_sizes.size(); ++split_idx) { in split_with_sizes_copy_out_cuda_contiguous_no_cast()
[all …]
|
| /aosp_15_r20/external/executorch/backends/vulkan/runtime/graph/ops/impl/ |
| H A D | Split.cpp | 23 const std::vector<int64_t>& split_sizes, in add_split_with_sizes_default_node() argument
34 VK_CHECK_COND(out_list->size() == split_sizes.size()); in add_split_with_sizes_default_node()
36 for (int split_idx = 0; split_idx < split_sizes.size(); split_idx++) { in add_split_with_sizes_default_node()
37 int64_t split_size = split_sizes[split_idx]; in add_split_with_sizes_default_node()
104 std::vector<int64_t> split_sizes = *(graph.get_int_list(split_sizes_ref)); in add_split_with_sizes_default_node() local
106 add_split_with_sizes_default_node(graph, in, split_sizes, dim, out); in add_split_with_sizes_default_node()
127 std::vector<int64_t> split_sizes(size / split_size, split_size); in add_split_tensor_node() local
129 add_split_with_sizes_default_node(graph, in, split_sizes, dim, out); in add_split_tensor_node()
|
| /aosp_15_r20/external/executorch/kernels/portable/cpu/ |
| H A D | op_split_with_sizes_copy.cpp | 26 exec_aten::ArrayRef<int64_t> split_sizes, in split_with_sizes_copy_out() argument
38 check_split_with_sizes_copy_args(in, split_sizes, dim, out), in split_with_sizes_copy_out()
59 for (size_t i = 0; i < split_sizes.size(); i++) { in split_with_sizes_copy_out()
60 target_out_sizes[dim] = static_cast<Tensor::SizesType>(split_sizes[i]); in split_with_sizes_copy_out()
87 size_t chunk_step = split_sizes[i] * trailing_dims; in split_with_sizes_copy_out()
90 target_out_sizes[dim] = static_cast<Tensor::SizesType>(split_sizes[i]); in split_with_sizes_copy_out()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| H A D | split_op.cc | 147 std::vector<int64_t> split_sizes; in Compile() local
148 OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &split_sizes)); in Compile()
151 int64_t slice_size = split_sizes[i]; in Compile()
181 split_sizes[neg_one_dim] = in Compile()
193 int slice_size = split_sizes[i]; in Compile()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/ |
| H A D | NestedTensorUtils.cpp | 115 c10::IntArrayRef split_sizes, in split_with_sizes_nested() argument
124 auto num_splits = split_sizes.size(); in split_with_sizes_nested()
131 for (const auto split_size : split_sizes) { in split_with_sizes_nested()
138 " (input tensor's size at dimension ", dim, "), but got split_sizes=", split_sizes); in split_with_sizes_nested()
151 auto split_size = split_sizes[split_idx]; in split_with_sizes_nested()
|
| H A D | NestedTensorMath.cpp | 276 std::vector<int64_t> split_sizes; in NestedTensor_to_padded_tensor_generic() local
277 split_sizes.reserve(sizes_num_rows); in NestedTensor_to_padded_tensor_generic()
280 split_sizes.push_back( in NestedTensor_to_padded_tensor_generic()
284 for (const auto split_size : split_sizes) { in NestedTensor_to_padded_tensor_generic()
296 buffers.reserve(split_sizes.size()); in NestedTensor_to_padded_tensor_generic()
299 for (const auto split_size : split_sizes) { in NestedTensor_to_padded_tensor_generic()
|
| /aosp_15_r20/external/executorch/kernels/test/ |
| H A D | op_split_with_sizes_copy_test.cpp | 24 exec_aten::ArrayRef<int64_t> split_sizes, in op_split_with_sizes_copy_out() argument
28 context_, self, split_sizes, dim, out); in op_split_with_sizes_copy_out()
43 exec_aten::ArrayRef<int64_t> split_sizes = exec_aten::ArrayRef<int64_t>( in test_tensor_shape_dynamism() local
103 op_split_with_sizes_copy_out(self, split_sizes, dim, out); in test_tensor_shape_dynamism()
|
| /aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/ |
| H A D | Utils.hpp | 499 const std::vector<int64_t>& split_sizes, in checkSplitSizes() argument
502 if (split_sizes.empty()) { in checkSplitSizes()
508 split_sizes.size() == static_cast<size_t>(group_size), in checkSplitSizes()
510 const auto sum = c10::sum_integers(split_sizes); in checkSplitSizes()
519 const std::vector<int64_t>& split_sizes, in computeLengthsAndOffsets() argument
530 if (split_sizes.empty()) { in computeLengthsAndOffsets()
535 size_t length = row_size * (equal_splits ? split_size : split_sizes[i]); in computeLengthsAndOffsets()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/ |
| H A D | FunctionalInverses.cpp | 254 …Mode inverse_return_mode, int64_t mutated_view_idx, c10::SymIntArrayRef split_sizes, int64_t dim) { in split_with_sizes_inverse() argument
259 start += split_sizes[i]; in split_with_sizes_inverse()
261 auto end = start + split_sizes[mutated_view_idx]; in split_with_sizes_inverse()
450 std::vector<c10::SymInt> split_sizes(chunks, split_size); in chunk_inverse() local
451 split_sizes[chunks - 1] = split_size - (split_size * chunks - dim_size); in chunk_inverse()
452 …it_with_sizes_inverse(base, mutated_view, inverse_return_mode, mutated_view_idx, split_sizes, dim); in chunk_inverse()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/functorch/ |
| H A D | LegacyBatchingRegistrations.cpp | 249 std::vector<Tensor> split_with_sizes_batching_rule(const Tensor& self, SymIntArrayRef split_sizes, … in split_with_sizes_batching_rule() argument
252 return split_with_sizes_symint(self, split_sizes, dim); in split_with_sizes_batching_rule()
256 auto result = split_with_sizes_symint(self_physical.tensor(), split_sizes, dim_physical); in split_with_sizes_batching_rule()
261 … split_with_sizes_copy_batching_rule(const Tensor& self, SymIntArrayRef split_sizes, int64_t dim) { in split_with_sizes_copy_batching_rule() argument
264 return split_with_sizes_copy_symint(self, split_sizes, dim); in split_with_sizes_copy_batching_rule()
268 auto result = split_with_sizes_copy_symint(self_physical.tensor(), split_sizes, dim_physical); in split_with_sizes_copy_batching_rule()
|
| /aosp_15_r20/external/executorch/backends/arm/test/ops/ |
| H A D | test_split.py | 37 def forward(self, x: torch.Tensor, split_sizes: list[int], dim: int):
38 return x.split_with_sizes(split_sizes=split_sizes, dim=dim)
|
| /aosp_15_r20/external/pytorch/test/distributed/ |
| H A D | test_functional_api.py | 523 split_sizes = [(i + 1) * (rank + 1) for i in range(self.world_size)]
525 x, output_split_sizes=split_sizes, input_split_sizes=split_sizes, group=mesh
528 for idx, tensor in enumerate(torch.split(x, split_sizes)):
541 split_sizes = [(i + 1) * (rank + 1) for i in range(self.world_size)]
543 x, output_split_sizes=split_sizes, input_split_sizes=split_sizes, group=mesh
546 for idx, tensor in enumerate(torch.split(x, split_sizes)):
|
| H A D | test_c10d_spawn.py | 240 split_sizes = [(i + 1) * (self.rank + 1) for i in range(self.world_size)]
242 y, x, output_split_sizes=split_sizes, input_split_sizes=split_sizes
245 for idx, tensor in enumerate(torch.split(x, split_sizes)):
|
| /aosp_15_r20/external/pytorch/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/ |
| H A D | embedding_bag.py | 422 split_sizes = torch.sum(
426 split_sizes = torch.cat(
433 return torch.div(result, split_sizes.unsqueeze(1))
|
| /aosp_15_r20/external/executorch/backends/apple/mps/operators/ |
| H A D | shape_ops.py | 243 split_sizes = eval_shape(cast(torch.SymInt, node.args[1]))
256 split_sizes=split_sizes,
|
| /aosp_15_r20/external/executorch/kernels/portable/cpu/util/ |
| H A D | copy_ops_util.cpp | 418 exec_aten::ArrayRef<int64_t> split_sizes, in check_split_with_sizes_copy_args() argument
425 split_sizes.size() == out.size(), in check_split_with_sizes_copy_args()
429 for (int i = 0; i < split_sizes.size(); i++) { in check_split_with_sizes_copy_args()
431 split_sizes[i] >= 0, "All split sizes must be non negative."); in check_split_with_sizes_copy_args()
432 sum += split_sizes[i]; in check_split_with_sizes_copy_args()
|
| H A D | copy_ops_util.h | 141 exec_aten::ArrayRef<int64_t> split_sizes,
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | TensorShape.cpp | 951 std::vector<c10::SymInt> split_sizes(chunks, split_size); in chunk() local
952 split_sizes[chunks - 1] = split_size - (split_size * chunks - dim_size); in chunk()
953 return self.split_with_sizes_symint(split_sizes, dim); in chunk()
1042 std::vector<int64_t> split_sizes(chunks, split_size); in unsafe_chunk() local
1043 split_sizes[chunks - 1] = split_size - (split_size * chunks - dim_size); in unsafe_chunk()
1044 return self.unsafe_split_with_sizes(split_sizes, dim); in unsafe_chunk()
2621 std::vector<Tensor> split_with_sizes(const Tensor& self, IntArrayRef split_sizes, int64_t dim) { in split_with_sizes() argument
2624 const int64_t num_splits = split_sizes.size(); in split_with_sizes()
2630 auto length = split_sizes[i]; in split_with_sizes()
2633 "entries, but got split_sizes=", split_sizes); in split_with_sizes()
[all …]
|
| /aosp_15_r20/external/pytorch/torch/_inductor/fx_passes/ |
| H A D | post_grad.py | 620 split_sizes = get_arg_value(split_node, 1, "split_sizes")
622 if get_item_args != set(range(len(split_sizes))):
630 if cat_items_args_order != list(range(len(split_sizes))):
937 split_sizes = get_arg_value(split_node, 1, "split_sizes")
940 if len(cat_inputs) != len(split_sizes):
943 for cat_input, split_size in zip(cat_inputs, split_sizes):
|
| /aosp_15_r20/external/pytorch/torch/onnx/ |
| H A D | symbolic_opset13.py | 73 split_sizes = [
83 "Add", start, split_sizes[i]
119 def split_with_sizes(g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None): argument
120 return split(g, self, split_sizes, dim, _outputs)
132 g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None argument
134 return split_with_sizes(g, self, split_sizes, dim, _outputs)
|
| H A D | symbolic_opset11.py | 610 split_sizes = [
619 "Add", start, split_sizes[i]
638 def split_with_sizes(g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None): argument
639 return split(g, self, split_sizes, dim, _outputs)
|
| /aosp_15_r20/external/pytorch/torch/_decomp/ |
| H A D | decompositions.py | 1398 self: Tensor, split_sizes: List[int], dim: int = 0
1402 for i in range(len(split_sizes)):
1404 split_sizes[i],
1409 sum(split_sizes) == self.shape[dim],
1412 num_splits = len(split_sizes)
1417 length = split_sizes[i]
1429 split_sizes: List[int],
1433 splits = split_with_sizes(self, split_sizes, dim=dim)
1450 input: Tensor, split_sizes: List[int], dim: int = 0
1452 return aten.split_with_sizes.default(input, split_sizes, dim)
[all …]
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | split_v_op.cc | 199 absl::Span<const Tlen> split_sizes) { in SplitHasAlignedOutputsInFirstDimension() argument
204 for (const Tlen split_size : split_sizes) { in SplitHasAlignedOutputsInFirstDimension()
|
| /aosp_15_r20/external/pytorch/torch/csrc/jit/runtime/static/ |
| H A D | native_ops.cpp | 733 const auto& split_sizes = p_node->Input(1).toIntList(); in __anon75e5f0514602() local
736 at::native::split_with_sizes(self, split_sizes.vec(), dim); in __anon75e5f0514602()
757 const auto& split_sizes = p_node->Input(1).toIntList(); in __anon75e5f0514902() local
760 at::native::split_with_sizes(self, split_sizes.vec(), dim); in __anon75e5f0514902()
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/ |
| H A D | cross_device_ops.py | 765 split_sizes = [split_size] * (num_splits - 1) + [split_size_last]
766 grad_packs = array_ops.split(concat_grads, split_sizes)
|