1 #include <ATen/core/Tensor.h>
2 #include <ATen/LegacyBatchedTensorImpl.h>
3 #include <ATen/WrapDimUtils.h>
4 #include <ATen/LegacyVmapTransforms.h>
5
6 #ifdef AT_PER_OPERATOR_HEADERS
7 #include <ATen/ops/_add_batch_dim_native.h>
8 #include <ATen/ops/_remove_batch_dim_native.h>
9 #endif
10
11 namespace at::native {
12
13 // Adds a batch dimension to the tensor `self` out-of-place
_add_batch_dim(const Tensor & self,int64_t batch_dim,int64_t level)14 Tensor _add_batch_dim(const Tensor& self, int64_t batch_dim, int64_t level) {
15 return addBatchDim(self, level, batch_dim);
16 }
17
has_level(const Tensor & self,int64_t level)18 static bool has_level(const Tensor& self, int64_t level) {
19 const auto* batched = maybeGetBatchedImpl(self);
20 if (!batched) {
21 return false;
22 }
23 auto bdims = batched->bdims();
24 auto* it = std::find_if(bdims.begin(), bdims.end(), [&](const BatchDim& bdim) {
25 return bdim.level() == level;
26 });
27 return it != bdims.end();
28 }
29
30 // Returns a Tensor with batch dim with level `level` turned into a regular dimension,
31 // as well as a logical dim index of where said dimension is in the returned tensor.
32 // A call to this function is always followed by a call to `movedim`.
33 //
34 // Preconditions: A BatchDim with level `level` must exist inside `batched`.
35 //
36 // The reason why we want to return the index of where said dimension is in the returned
37 // tensor is because we want to keep track of which dimension used to be the batch
38 // dimension so that we can move it to the correct logical dimension specified by
39 // `out_dims` in vmap. For example, if we had
40 // >>> x = torch.randn(2, 3, 5)
41 // >>> vmap(lambda x: x, in_dims=0, out_dims=1)(x)
42 // then right when we are about to exit the vmap block, x is a BatchedTensor with a
43 // batch dimension at (physical) index 0. Note that the batch dimension doesn't
44 // always have to exist at (physical) index 0. When we undo the batch dimension,
45 // we want to move it to dimension 1 (as specified by out_dims). So we return the
46 // index at which the batch dim appears so that we can move it to the correct place.
47 // later down the line via a call to `movedim`.
remove_existing_batch_dim(const BatchedTensorImpl * batched,int64_t level)48 static std::pair<Tensor,int64_t> remove_existing_batch_dim(
49 const BatchedTensorImpl* batched, int64_t level) {
50 auto bdims = batched->bdims();
51 if (bdims.size() == 1) {
52 TORCH_INTERNAL_ASSERT(bdims[0].level() == level);
53 return std::make_pair(batched->value(), bdims[0].dim());
54 }
55 BatchDims new_bdims;
56 int64_t newly_exposed_physical_dim = -1;
57 new_bdims.reserve(bdims.size() - 1);
58 for (const auto& bdim : bdims) {
59 if (bdim.level() == level) {
60 newly_exposed_physical_dim = bdim.dim();
61 } else {
62 new_bdims.push_back(bdim);
63 }
64 }
65 // Because a BatchDim with level `level` must exist inside `batched,
66 // we should have found a `newly_exposed_logical_dim`.
67 TORCH_INTERNAL_ASSERT(newly_exposed_physical_dim != -1);
68 int64_t num_batch_dims_before_newly_exposed_physical_dim = std::count_if(
69 new_bdims.begin(), new_bdims.end(),
70 [&](const BatchDim& bdim) {
71 return bdim.dim() < newly_exposed_physical_dim;
72 });
73 int64_t newly_exposed_logical_dim =
74 newly_exposed_physical_dim - num_batch_dims_before_newly_exposed_physical_dim;
75 auto result_tensor = makeBatched(batched->value(), std::move(new_bdims));
76 return std::make_pair(std::move(result_tensor), newly_exposed_logical_dim);
77 }
78
79 // at::movedim but may return the original tensor if dst is the same as src.
maybe_movedim(const Tensor & self,int64_t src,int64_t dst)80 static Tensor maybe_movedim(const Tensor& self, int64_t src, int64_t dst) {
81 auto logical_dim = self.dim();
82 src = maybe_wrap_dim(src, logical_dim);
83 dst = maybe_wrap_dim(dst, logical_dim);
84 if (src == dst) {
85 return self;
86 }
87 return self.movedim(src, dst);
88 }
89
90 // Removes the batch dim with level `level` from `self`. If this causes the
91 // last batch dim to be removed from a BatchedTensor, then this returns a
92 // regular Tensor.
93 //
94 // If the `level` of the batch dim to remove does not exist in `self`, then we
95 // add the batch dim in. This can happen if `self` didn't interact with a tensor
96 // inside the vmap level, for example,
97 // self = torch.randn(3)
98 // y = torch.randn(5)
99 // out = vmap(lambda x: vmap(lambda y: x)(y))(self)
100 // assert out.shape == (3, 5)
101 // Inside the inner vmap, `x` is a BatchedTensor with a single batch dimension
102 // corresponding to the *outer* vmap level and it doesn't have any dimensions that
103 // correspond to the inner vmap level so we need to create one for the user.
104 //
105 // `out_dim` controls where we should put the batch dimension in the output tensor.
_remove_batch_dim(const Tensor & self,int64_t level,int64_t batch_size,int64_t out_dim)106 Tensor _remove_batch_dim(const Tensor& self, int64_t level, int64_t batch_size, int64_t out_dim) {
107 if (!has_level(self, level)) {
108 auto self_sizes = self.sizes();
109 VmapDimVector expanded_sizes(self_sizes.begin(), self_sizes.end());
110 expanded_sizes.insert(expanded_sizes.begin() + out_dim, batch_size);
111 return self.expand(expanded_sizes);
112 }
113
114 // Must be batched if has_level(self, /*any_level*/)
115 const auto* batched = maybeGetBatchedImpl(self);
116 TORCH_INTERNAL_ASSERT(batched != nullptr);
117
118 auto [self_without_bdim, newly_exposed_logical_dim] = remove_existing_batch_dim(batched, level);
119 return maybe_movedim(self_without_bdim, newly_exposed_logical_dim, out_dim);
120 }
121
122 } // namespace at::native
123