| /external/pytorch/aten/src/ATen/ |
| D | LegacyVmapTransforms.cpp | 27 VmapDimVector permutation(sizes.size(), 0); in permuteBatchDimsToFront()
59 VmapDimVector VmapPhysicalView::getPhysicalDims(OptionalIntArrayRef opt_logical_dims) const { in getPhysicalDims()
62 VmapDimVector result; in getPhysicalDims()
82 VmapDimVector VmapPhysicalView::getPhysicalShape(IntArrayRef logical_shape) const { in getPhysicalShape()
83 VmapDimVector result; in getPhysicalShape()
158 VmapDimVector aligned_sizes(requested_levels.count() + requested_example_dim, 1); in alignBatchDimsAtFront()
214 VmapDimVector batch_sizes(num_batch_dims, 1); in logicalToPhysical()
227 VmapDimVector expanded_size(batch_sizes.begin(), batch_sizes.end()); in logicalToPhysical()
|
| D | LegacyVmapTransforms.h | 42 using VmapDimVector = SmallVector<int64_t, kVmapStaticDimVecSize>; variable
136 VmapDimVector getPhysicalDims(OptionalIntArrayRef logical_dims) const;
145 VmapDimVector getPhysicalShape(IntArrayRef logical_shape) const;
|
| D | LegacyBatchedFallback.cpp | 112 VmapDimVector batched_tensor_inputs_position; in batchedTensorInplaceForLoopFallback()
276 VmapDimVector batched_tensor_inputs_position; in batchedTensorForLoopFallback()
367 VmapDimVector output_sizes(batch_sizes); in batchedTensorForLoopFallback()
|
| D | LegacyBatchingRegistrations.cpp | 182 VmapDimVector view_shape(size_physical.size(), 1); in expand_batching_rule()
278 VmapDimVector squeezed_sizes; in squeeze_batching_rule()
348 VmapDimVector all_dims_physical; in permute_batching_rule()
641 at::VmapDimVector physical_strides; in as_strided_batching_rule()
|
| /external/pytorch/aten/src/ATen/functorch/ |
| D | LegacyVmapTransforms.h | 49 using VmapDimVector = SmallVector<int64_t, kVmapStaticDimVecSize>; variable
139 VmapDimVector getPhysicalDims(IntArrayRef logical_dims) const;
148 VmapDimVector getPhysicalShape(IntArrayRef logical_shape) const;
|
| D | LegacyVmapTransforms.cpp | 23 VmapDimVector permutation(sizes.size(), 0); in permuteBatchDimsToFront()
53 VmapDimVector VmapPhysicalView::getPhysicalDims(IntArrayRef logical_dims) const { in getPhysicalDims()
56 VmapDimVector result; in getPhysicalDims()
69 VmapDimVector VmapPhysicalView::getPhysicalShape(IntArrayRef logical_shape) const { in getPhysicalShape()
70 VmapDimVector result; in getPhysicalShape()
|
| D | BatchRulesHelper.cpp | 55 VmapDimVector getPhysicalDims(const Tensor& tensor, bool has_batch_dim, IntArrayRef logical_dims) { in getPhysicalDims()
59 VmapDimVector result; in getPhysicalDims()
108 VmapDimVector new_shape(x.sizes().begin(), x.sizes().end()); in reshape_dim_into()
116 VmapDimVector shape(x.sizes().begin(), x.sizes().end()); in reshape_dim_outof()
|
| D | BatchedFallback.cpp | 117 VmapDimVector batched_tensor_inputs_position; in batchedTensorInplaceForLoopFallback()
292 VmapDimVector batched_tensor_inputs_position; in batchedTensorForLoopFallback()
388 VmapDimVector output_sizes(batch_sizes); in batchedTensorForLoopFallback()
433 VmapDimVector batched_tensor_inputs_position; in batchedNestedTensorForLoopFallback()
|
| D | BatchRulesViews.cpp | 142 VmapDimVector new_dims; in flip_batch_rule()
176 VmapDimVector new_size(size); in resize__plumbing()
286 VmapDimVector new_dims; in roll_batch_rule()
382 VmapDimVector dims_; in permute_batching_rule()
|
| D | BatchRulesScatterOps.cpp | 165 VmapDimVector permutation(tensor.dim(), 0); in swap_regions()
349 VmapDimVector compute_indexed_shape(const Tensor &src, TensorList indices_list) in compute_indexed_shape()
368 auto shape = VmapDimVector(src.sizes()); in compute_indexed_shape()
378 VmapDimVector get_indexed_shape(Tensor self, const torch::List<std::optional<at::Tensor>> &orig) in get_indexed_shape()
434 VmapDimVector new_values_shape(values_sizes.size() + n_unit_dims);
833 VmapDimVector new_index_shape(self_size.size(), 1); in get_expanded_index()
839 VmapDimVector new_index_shape = {self_size.begin(), self_size.end()}; in get_expanded_index()
|
| D | BatchRulesHelper.h | 38 VmapDimVector getPhysicalDims(const Tensor& tensor, bool has_batch_dim, IntArrayRef logical_dims);
462 inline VmapDimVector range(int64_t start, int64_t stop) {
464 VmapDimVector dims;
|
| D | BatchRulesFactory.cpp | 42 VmapDimVector new_shape; in apply()
|
| D | BatchRulesReduceOps.cpp | 166 VmapDimVector new_dims; in boxed_reduction_batch_rule()
|
| D | BatchRulesNorm.cpp | 37 VmapDimVector new_sizes(tensor.sizes().begin(), tensor.sizes().end()); in padRight()
|
| /external/pytorch/aten/src/ATen/test/ |
| D | legacy_vmap_test.cpp | 297 VmapDimVector({3, 4, 4, 3})); in TEST()
671 ASSERT_EQ(physical_view.getPhysicalShape({}), VmapDimVector({2, 3})); in TEST()
672 ASSERT_EQ(physical_view.getPhysicalShape({7}), VmapDimVector({2, 3, 7})); in TEST()
673 ASSERT_EQ(physical_view.getPhysicalShape({7, 11, 13}), VmapDimVector({2, 3, 7, 11, 13})); in TEST()
674 … ASSERT_EQ(physical_view.getPhysicalShape({7, 11, 13, 17}), VmapDimVector({2, 3, 7, 11, 13, 17})); in TEST()
678 ASSERT_EQ(physical_view.getPhysicalShape({}), VmapDimVector({2})); in TEST()
679 ASSERT_EQ(physical_view.getPhysicalShape({7}), VmapDimVector({2, 7})); in TEST()
|
| /external/pytorch/aten/src/ATen/native/ |
| D | LegacyBatching.cpp | 109 VmapDimVector expanded_sizes(self_sizes.begin(), self_sizes.end()); in _remove_batch_dim()
|
| /external/pytorch/torch/csrc/functorch/ |
| D | init.cpp | 113 VmapDimVector permutation; in _movedim()
154 VmapDimVector expanded_sizes(self_sizes.begin(), self_sizes.end()); in _remove_batch_dim()
|