1 #include <torch/csrc/python_headers.h> 3 #include <torch/csrc/utils/tensor_layouts.h> 5 #include <torch/csrc/Layout.h> 6 #include <torch/csrc/DynamicTypes.h> 7 #include <torch/csrc/Exceptions.h> 9 #include <c10/core/ScalarType.h> 10 #include <ATen/Layout.h> 12 namespace torch {
namespace utils {
14 void initializeLayouts() {
15 auto torch_module =
THPObjectPtr(PyImport_ImportModule(
"torch"));
18 PyObject *strided_layout = THPLayout_New(at::Layout::Strided,
"torch.strided");
19 Py_INCREF(strided_layout);
20 if (PyModule_AddObject(torch_module,
"strided", strided_layout) != 0) {
24 registerLayoutObject((
THPLayout*)strided_layout, at::Backend::CPU);
25 registerLayoutObject((
THPLayout*)strided_layout, at::Backend::CUDA);
26 registerLayoutObject((
THPLayout*)strided_layout, at::Backend::MSNPU);
27 registerLayoutObject((
THPLayout*)strided_layout, at::Backend::XLA);
29 PyObject *sparse_coo_layout = THPLayout_New(at::Layout::Sparse,
"torch.sparse_coo");
30 Py_INCREF(sparse_coo_layout);
31 if (PyModule_AddObject(torch_module,
"sparse_coo", sparse_coo_layout) != 0) {
34 registerLayoutObject((
THPLayout*)sparse_coo_layout, at::Backend::SparseCPU);
35 registerLayoutObject((
THPLayout*)sparse_coo_layout, at::Backend::SparseCUDA);