Home
last modified time | relevance | path

Searched refs:VmapPhysicalView (Results 1 – 7 of 7) sorted by relevance

/external/pytorch/aten/src/ATen/
DLegacyVmapTransforms.h31 struct VmapPhysicalView;
36 SmallVector<VmapPhysicalView, kVmapTransformStaticInputSize>;
59 static VmapPhysicalView logicalToPhysical(const Tensor& logical_tensor);
114 struct TORCH_API VmapPhysicalView { struct
115 VmapPhysicalView(Tensor&& tensor, std::bitset<kVmapNumLevels> levels) in VmapPhysicalView() function
DLegacyVmapTransforms.cpp43 VmapPhysicalView MultiBatchVmapTransform::logicalToPhysical(const Tensor& logical_tensor) { in logicalToPhysical()
51 int64_t VmapPhysicalView::numBatchDims() const { in numBatchDims()
55 int64_t VmapPhysicalView::numLogicalDims() const { in numLogicalDims()
59 VmapDimVector VmapPhysicalView::getPhysicalDims(OptionalIntArrayRef opt_logical_dims) const { in getPhysicalDims()
77 int64_t VmapPhysicalView::getPhysicalDim(int64_t logical_dim) const { in getPhysicalDim()
82 VmapDimVector VmapPhysicalView::getPhysicalShape(IntArrayRef logical_shape) const { in getPhysicalShape()
282 VmapPhysicalToLogicalMap VmapPhysicalView::getPhysicalToLogicalMap() const { in getPhysicalToLogicalMap()
DLegacyBatchingRegistrations.cpp934 physical_views, [](const VmapPhysicalView& view) -> Tensor { return view.tensor(); }); in cat_batching_rule()
944 physical_views, [](const VmapPhysicalView& view) -> Tensor { return view.tensor(); }); in stack_batching_rule()
/external/pytorch/aten/src/ATen/functorch/
DLegacyVmapTransforms.h39 struct VmapPhysicalView;
43 using VmapPhysicalViewVec = SmallVector<VmapPhysicalView, kVmapTransformStaticInputSize>;
66 static VmapPhysicalView logicalToPhysical(const Tensor& logical_tensor);
121 struct TORCH_API VmapPhysicalView { struct
122 VmapPhysicalView(Tensor&& tensor, std::bitset<kVmapNumLevels> levels) in VmapPhysicalView() argument
DLegacyVmapTransforms.cpp37 VmapPhysicalView MultiBatchVmapTransform::logicalToPhysical(const Tensor& logical_tensor) { in logicalToPhysical()
45 int64_t VmapPhysicalView::numBatchDims() const { in numBatchDims()
49 int64_t VmapPhysicalView::numLogicalDims() const { in numLogicalDims()
53 VmapDimVector VmapPhysicalView::getPhysicalDims(IntArrayRef logical_dims) const { in getPhysicalDims()
64 int64_t VmapPhysicalView::getPhysicalDim(int64_t logical_dim) const { in getPhysicalDim()
69 VmapDimVector VmapPhysicalView::getPhysicalShape(IntArrayRef logical_shape) const { in getPhysicalShape()
78 SymDimVector VmapPhysicalView::getPhysicalShape(c10::SymIntArrayRef logical_shape) const { in getPhysicalShape()
214 VmapPhysicalToLogicalMap VmapPhysicalView::getPhysicalToLogicalMap() const { in getPhysicalToLogicalMap()
DLegacyBatchingRegistrations.cpp588 physical_views, [](const VmapPhysicalView& view) -> Tensor { return view.tensor(); }); in block_diag_batching_rule()
616 physical_views, [](const VmapPhysicalView& view) -> Tensor { return view.tensor(); }); in stack_batching_rule()
/external/pytorch/aten/src/ATen/test/
Dlegacy_vmap_test.cpp276 VmapPhysicalView physical_view(ones({2, 3, 4, 5, 6}), 1 | 4); in TEST()
293 VmapPhysicalView physical_view(ones({2, 3, 4, 5, 6}), 2 | 8 | 16); in TEST()
316 VmapPhysicalView physical_view(ones({2, 3, 4}), /*levels = {2}*/4); in TEST()
327 VmapPhysicalView physical_view(ones({2, 3, 4, 5, 6}), /*levels = {1, 3, 4}*/2 | 8 | 16); in TEST()
338 VmapPhysicalView physical_view(ones({2}), /*levels = {2}*/4); in TEST()
670 VmapPhysicalView physical_view(ones({2, 3, 4, 5, 6}), 1 | 4); in TEST()
677 VmapPhysicalView physical_view(ones({2, 3, 4, 5, 6}), 2); in TEST()