forked from ROCm/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ExpandBase.h
30 lines (26 loc) · 914 Bytes
/
ExpandBase.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
#include <ATen/core/TensorBase.h>
// Broadcasting utilities for working with TensorBase
namespace at {
namespace internal {
TORCH_API TensorBase expand_slow_path(const TensorBase& self, IntArrayRef size);
} // namespace internal
inline c10::MaybeOwned<TensorBase> expand_size(
const TensorBase& self,
IntArrayRef size) {
if (size.equals(self.sizes())) {
return c10::MaybeOwned<TensorBase>::borrowed(self);
}
return c10::MaybeOwned<TensorBase>::owned(
at::internal::expand_slow_path(self, size));
}
c10::MaybeOwned<TensorBase> expand_size(TensorBase&& self, IntArrayRef size) =
delete;
inline c10::MaybeOwned<TensorBase> expand_inplace(
const TensorBase& tensor,
const TensorBase& to_expand) {
return expand_size(to_expand, tensor.sizes());
}
c10::MaybeOwned<TensorBase> expand_inplace(
const TensorBase& tensor,
TensorBase&& to_expand) = delete;
} // namespace at