Home
last modified time | relevance | path

Searched refs:tensor_split (Results 1 – 25 of 33) sorted by relevance

12

/aosp_15_r20/external/pytorch/test/
H A Dtest_view_ops.py400 a_split_dim0 = a.tensor_split(7, 0)
403 a_split_dim1 = a.tensor_split(7, 1)
1839 result1 = torch.tensor_split(a, sections, dim)
1840 result2 = torch.tensor_split(
1883 result_1 = torch.tensor_split(a, indices, dim)
1884 result_2 = torch.tensor_split(
1945 torch.tensor_split(a, sections_or_indices, dim)
1947 torch.tensor_split(a, torch.tensor(sections_or_indices), dim)
1956 torch.tensor_split(a, torch.tensor(1.1), dim)
1963 torch.tensor_split(torch.rand(S, device=device), torch.tensor(((1,),)), 0)
/aosp_15_r20/external/pytorch/torch/distributed/pipelining/
H A Dmicrobatch.py190 chunk_tensors = torch.tensor_split(
430 meta_chunks = torch.tensor_split(
H A Dschedules.py610 targets_split = list(torch.tensor_split(target, self._n_microbatches))
1199 targets_split = list(torch.tensor_split(target, self._n_microbatches))
/aosp_15_r20/external/pytorch/test/mobile/model_test/
H A Dtensor_ops.py206 torch.tensor_split(x, 1),
207 torch.tensor_split(x, [0, 1]),
/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DNonSymbolicBC.h25 TORCH_API std::vector<Tensor> tensor_split(const Tensor& self, IntArrayRef indices, int64_t dim);
H A DTensorShape.cpp995 std::vector<Tensor> tensor_split(const Tensor& self, IntArrayRef indices, int64_t dim) { in tensor_split() function
1003 std::vector<Tensor> tensor_split(const Tensor& self, const Tensor& tensor_indices_or_sections, int6… in tensor_split() function
1017 return self.tensor_split(sections, dim); in tensor_split()
1027 return self.tensor_split(indices, dim); in tensor_split()
2604 return at::tensor_split(self, split_size, dim); in hsplit()
2611 return at::tensor_split(self, split_size, 0); in vsplit()
2618 return at::tensor_split(self, split_size, 2); in dsplit()
2656 return at::tensor_split(self, split_sizes, (self.dim() == 1) ? 0 : 1); in hsplit()
2661 return at::tensor_split(self, split_sizes, 0); in vsplit()
2666 return at::tensor_split(self, split_sizes, 2); in dsplit()
/aosp_15_r20/external/pytorch/test/export/
H A Dtest_pass_infra.py20 return torch.ops.aten.tensor_split.sections(y, 2)
/aosp_15_r20/external/pytorch/torch/csrc/jit/runtime/static/
H A Dnative_ops.cpp860 REGISTER_NATIVE_OPERATOR_FUNCTOR(aten::tensor_split, aten_tensor_split, [](Node* n) -> SROperator { in __anon75e5f0515402()
867 pnode->Output(0) = at::native::tensor_split(a, b, c); in __anon75e5f0515402()
887 pnode->Output(0) = at::native::tensor_split(a, b, c); in __anon75e5f0515402()
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcustom_op_db.py316 return [xi.clone() for xi in torch.tensor_split(x, splits, dim)]
349 return [xi.clone() for xi in torch.tensor_split(x, splits, dim)], len(splits)
/aosp_15_r20/external/pytorch/docs/source/
H A Dtensor_view.rst80 - :meth:`~torch.Tensor.tensor_split`
H A Dtorch.rst133 tensor_split
H A Dtensors.rst729 Tensor.tensor_split
/aosp_15_r20/external/pytorch/torch/_functorch/
H A Dvmap.py360 t.tensor_split(split_idxs, dim=in_dim)
/aosp_15_r20/external/executorch/backends/apple/mps/test/
H A Dtest_mps_models.py154 split = torch.ops.aten.tensor_split.sections(x, 3)
/aosp_15_r20/external/pytorch/functorch/op_analysis/
H A Dpublic_api208 tensor_split
H A Dannotated_ops70 tensor_split, view/reshape
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DLegacyBatchingRegistrations.cpp222 auto result = at::tensor_split(self_physical.tensor(), sections, dim_physical); in tensor_split_sections_batching_rule()
230 auto result = at::tensor_split(self_physical.tensor(), indices, dim_physical); in tensor_split_indices_batching_rule()
/aosp_15_r20/external/executorch/exir/tests/
H A Dtest_memory_planning.py134 s0, s1 = torch.tensor_split(a, 2)
H A Dmodels.py222 return torch.tensor_split(input, sections, dim)
H A Dtest_passes.py462 return torch.ops.aten.tensor_split.sections(y, 2)
494 return torch.ops.aten.tensor_split.sections(y, 2)
/aosp_15_r20/external/pytorch/torch/_refs/
H A D__init__.py3956 torch.squeeze(s, dim) for s in torch.tensor_split(t, t.shape[dim], dim)
4106 def tensor_split( function
4206 return tensor_split(a, split_size, dim)
4218 return tensor_split(a, split_sizes, dim)
4246 return tensor_split(a, split_size, 0)
4258 return tensor_split(a, split_sizes, 0)
4455 return tensor_split(a, sections, 2)
/aosp_15_r20/external/executorch/test/end2end/
H A Dtest_end2end.py206 split = torch.ops.aten.tensor_split.sections(x, 3)
/aosp_15_r20/external/pytorch/torch/_refs/nn/functional/
H A D__init__.py1180 b, c = torch.tensor_split(a, 2, dim)
/aosp_15_r20/external/pytorch/torch/
H A D_torch_docs.py1834 torch.tensor_split,
/aosp_15_r20/external/pytorch/torch/onnx/
H A Dsymbolic_opset13.py139 def tensor_split( function

12