forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
THTensor.cpp
72 lines (58 loc) · 2.33 KB
/
THTensor.cpp
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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
#include <TH/THTensor.hpp>
#include <TH/generic/THTensor.cpp>
#include <TH/THGenerateAllTypes.h>
#include <TH/generic/THTensor.cpp>
#include <TH/THGenerateComplexTypes.h>
#include <TH/generic/THTensor.cpp>
#include <TH/THGenerateHalfType.h>
#include <TH/generic/THTensor.cpp>
#include <TH/THGenerateBoolType.h>
#include <TH/generic/THTensor.cpp>
#include <TH/THGenerateBFloat16Type.h>
#include <ATen/native/Resize.h>
#include <ATen/TensorUtils.h>
#include <numeric>
// NB: This is NOT valid on UndefinedTensorImpl
void THTensor_free(THTensor *self)
{
if (!self) return;
c10::raw::intrusive_ptr::decref(self);
}
void THTensor_setStorage(THTensor *self, THStorage *storage_, ptrdiff_t storageOffset_, at::IntArrayRef size_, at::IntArrayRef stride_) {
c10::raw::intrusive_ptr::incref(storage_);
THTensor_wrap(self).set_(at::Storage(c10::intrusive_ptr<at::StorageImpl>::reclaim(storage_)), storageOffset_, size_, stride_);
}
void THTensor_resize(THTensor *self, at::IntArrayRef size, at::IntArrayRef stride)
{
if (stride.data()) {
THArgCheck(stride.size() == size.size(), 3, "invalid stride");
}
#ifdef DEBUG
THAssert(size.size() <= INT_MAX);
#endif
THTensor_resizeNd(self, size.size(), size.data(), stride.data());
}
void THTensor_resizeNd(THTensor *self, int nDimension, const int64_t *size, const int64_t *stride)
{
TORCH_CHECK(nDimension >= 0, "resizeNd nDimension must be non-negative");
at::IntArrayRef sizes(size, nDimension);
at::optional<at::IntArrayRef> strides;
if (stride) {
strides = at::IntArrayRef(stride, nDimension);
}
at::native::resize_impl_cpu_(self, sizes, strides);
}
// NB: Steals ownership of storage
void THTensor_stealAndSetStoragePtr(THTensor* tensor, THStorage* storage) {
// Caffe2 might have tensors whose storages are null, but we
// don't allow it in PyTorch.
AT_ASSERT(storage);
// We used to allow this, but this breaks device caching.
// Let's put an actual error message for this one.
TORCH_CHECK(tensor->storage().device() == storage->device(),
"Attempted to set the storage of a tensor on device \"", tensor->storage().device(),
"\" to a storage on different device \"", storage->device(),
"\". This is no longer allowed; the devices must match.");
tensor->set_storage_keep_dtype(
at::Storage(c10::intrusive_ptr<THStorage>::reclaim(storage)));
}