Validate Sparse Compressed tensor inputs (pytorch#79385)

The validation includes regular tensor inputs, batched tensor inputs, as well as hybrid tensor inputs. Pull Request resolved: pytorch#79385 Approved by: https://github.com/nikitaved, https://github.com/cpuhrsch
walterddr · Jun 27, 2022 · cde365a · cde365a
1 parent 9db3c51
commit cde365a
Show file tree

Hide file tree

Showing 12 changed files with 4,157 additions and 2,616 deletions.
diff --git a/aten/src/ATen/SparseCsrTensorUtils.h b/aten/src/ATen/SparseCsrTensorUtils.h
@@ -181,6 +181,38 @@ inline std::string plainIndicesName(Layout layout) {
       [&] { return "row_indices"; });
 }
 
+inline std::string compressedDimName(Layout layout) {
+  switch (layout) {
+    case kSparseCsr:
+      return "row";
+    case kSparseCsc:
+      return "column";
+    case kSparseBsr:
+      return "row block";
+    case kSparseBsc:
+      return "column block";
+    default:
+      TORCH_CHECK(false, "Not a sparse compressed layout:", layout);
+      return "";
+  }
+}
+
+inline std::string plainDimName(Layout layout) {
+  switch (layout) {
+    case kSparseCsr:
+      return "column";
+    case kSparseCsc:
+      return "row";
+    case kSparseBsr:
+      return "column block";
+    case kSparseBsc:
+      return "row block";
+    default:
+      TORCH_CHECK(false, "Not a sparse compressed layout:", layout);
+      return "";
+  }
+}
+
 inline int rowDimension(Layout layout, IntArrayRef size) {
   return size.size() - (isCompressedRow(layout) ? 2 : 1);
 }

diff --git a/aten/src/ATen/native/sparse/SparseCsrTensor.cpp b/aten/src/ATen/native/sparse/SparseCsrTensor.cpp