28 #include <torch/script.h>
31 #include <type_traits>
36 #define CHECK_CUDA(x) \
38 TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") \
41 #define CHECK_CONTIGUOUS(x) \
43 TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") \
46 #define CHECK_TYPE(x, type) \
48 TORCH_CHECK(x.dtype() == torch::type, #x " must have type " #type) \
51 #define CHECK_SAME_DEVICE_TYPE(...) \
53 if (!SameDeviceType({__VA_ARGS__})) { \
57 " must all have the same device type but got " + \
58 TensorInfoStr({__VA_ARGS__})) \
62 #define CHECK_SAME_DTYPE(...) \
64 if (!SameDtype({__VA_ARGS__})) { \
67 " must all have the same dtype but got " + \
68 TensorInfoStr({__VA_ARGS__})) \
76 TORCH_CHECK(
false,
"Unsupported type");
100 return torch::kFloat32;
104 return torch::kFloat64;
108 template <
class T,
class TDtype>
110 return ToTorchDtype<T>() == t;
115 if (tensors.size()) {
116 auto device_type = tensors.begin()->device().type();
117 for (
auto t : tensors) {
118 if (device_type != t.device().type()) {
127 inline bool SameDtype(std::initializer_list<torch::Tensor> tensors) {
128 if (tensors.size()) {
129 auto dtype = tensors.begin()->dtype();
130 for (
auto t : tensors) {
131 if (dtype != t.dtype()) {
139 inline std::string
TensorInfoStr(std::initializer_list<torch::Tensor> tensors) {
140 std::stringstream sstr;
142 for (
const auto t : tensors) {
143 sstr << t.sizes() <<
" " << t.toString() <<
" " << t.device();
145 if (
count < tensors.size()) sstr <<
", ";
152 const torch::Device& device,
153 void** ptr =
nullptr) {
154 torch::Tensor tensor = torch::empty(
157 *ptr = tensor.data_ptr<uint8_t>();
163 torch::Tensor tensor) {
166 std::vector<DimValue> shape;
167 const int rank = tensor.dim();
168 for (
int i = 0; i < rank; ++i) {
169 shape.push_back(tensor.size(i));
177 std::tuple<bool, std::string>
CheckShape(torch::Tensor tensor,
180 return open3d::ml::op_util::CheckShape<Opt>(
GetShapeVector(tensor),
181 std::forward<TDimX>(dimex),
182 std::forward<TArgs>(args)...);
205 #define CHECK_SHAPE(tensor, ...) \
208 std::string cs_errstr_; \
209 std::tie(cs_success_, cs_errstr_) = CheckShape(tensor, __VA_ARGS__); \
210 TORCH_CHECK(cs_success_, \
211 "invalid shape for '" #tensor "', " + cs_errstr_) \
214 #define CHECK_SHAPE_COMBINE_FIRST_DIMS(tensor, ...) \
217 std::string cs_errstr_; \
218 std::tie(cs_success_, cs_errstr_) = \
219 CheckShape<CSOpt::COMBINE_FIRST_DIMS>(tensor, __VA_ARGS__); \
220 TORCH_CHECK(cs_success_, \
221 "invalid shape for '" #tensor "', " + cs_errstr_) \
224 #define CHECK_SHAPE_IGNORE_FIRST_DIMS(tensor, ...) \
227 std::string cs_errstr_; \
228 std::tie(cs_success_, cs_errstr_) = \
229 CheckShape<CSOpt::IGNORE_FIRST_DIMS>(tensor, __VA_ARGS__); \
230 TORCH_CHECK(cs_success_, \
231 "invalid shape for '" #tensor "', " + cs_errstr_) \
234 #define CHECK_SHAPE_COMBINE_LAST_DIMS(tensor, ...) \
237 std::string cs_errstr_; \
238 std::tie(cs_success_, cs_errstr_) = \
239 CheckShape<CSOpt::COMBINE_LAST_DIMS>(tensor, __VA_ARGS__); \
240 TORCH_CHECK(cs_success_, \
241 "invalid shape for '" #tensor "', " + cs_errstr_) \
244 #define CHECK_SHAPE_IGNORE_LAST_DIMS(tensor, ...) \
247 std::string cs_errstr_; \
248 std::tie(cs_success_, cs_errstr_) = \
249 CheckShape<CSOpt::IGNORE_LAST_DIMS>(tensor, __VA_ARGS__); \
250 TORCH_CHECK(cs_success_, \
251 "invalid shape for '" #tensor "', " + cs_errstr_) \
std::tuple< bool, std::string > CheckShape(torch::Tensor tensor, TDimX &&dimex, TArgs &&... args)
Definition: TorchHelper.h:177
TorchDtype_t ToTorchDtype< int64_t >()
Definition: TorchHelper.h:95
TorchDtype_t ToTorchDtype< uint8_t >()
Definition: TorchHelper.h:79
std::string TensorInfoStr(std::initializer_list< torch::Tensor > tensors)
Definition: TorchHelper.h:139
std::vector< open3d::ml::op_util::DimValue > GetShapeVector(torch::Tensor tensor)
Definition: TorchHelper.h:162
TorchDtype_t ToTorchDtype< int16_t >()
Definition: TorchHelper.h:87
TorchDtype_t ToTorchDtype< int8_t >()
Definition: TorchHelper.h:83
TorchDtype_t ToTorchDtype< double >()
Definition: TorchHelper.h:103
bool SameDtype(std::initializer_list< torch::Tensor > tensors)
Definition: TorchHelper.h:127
bool SameDeviceType(std::initializer_list< torch::Tensor > tensors)
Definition: TorchHelper.h:114
std::remove_const< decltype(torch::kInt32)>::type TorchDtype_t
Definition: TorchHelper.h:73
TorchDtype_t ToTorchDtype()
Definition: TorchHelper.h:75
torch::Tensor CreateTempTensor(const int64_t size, const torch::Device &device, void **ptr=nullptr)
Definition: TorchHelper.h:151
TorchDtype_t ToTorchDtype< int32_t >()
Definition: TorchHelper.h:91
bool CompareTorchDtype(const TDtype &t)
Definition: TorchHelper.h:109
TorchDtype_t ToTorchDtype< float >()
Definition: TorchHelper.h:99
Definition: ShapeChecking.h:35
CSOpt
Check shape options.
Definition: ShapeChecking.h:424