[tvm] branch unity updated: [Unity][Fix] Allow scalar layout initialization (#14370)

[ruihangl]

This is an automated email from the ASF dual-hosted git repository.

ruihangl pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new b8eb779ff3 [Unity][Fix] Allow scalar layout initialization (#14370)
b8eb779ff3 is described below

commit b8eb779ff3f1819ade4a4a68d2b9afc03585892e
Author: Prakalp Srivastava <prak...@octoml.ai>
AuthorDate: Wed Mar 22 17:11:28 2023 -0400

    [Unity][Fix] Allow scalar layout initialization (#14370)
    
    Fix assert to allow scalar (ndim=0) layout initialization.
---
 src/relax/transform/infer_layout_utils.cc           |  4 ++--
 tests/python/relax/test_transform_convert_layout.py | 13 +++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/relax/transform/infer_layout_utils.cc 
b/src/relax/transform/infer_layout_utils.cc
index e603fb4a1b..d746f9394a 100644
--- a/src/relax/transform/infer_layout_utils.cc
+++ b/src/relax/transform/infer_layout_utils.cc
@@ -53,7 +53,7 @@ int FindAxis(const Layout& dst, int axis) {
 }
 
 Layout InitialLayout(int ndim) {
-  ICHECK(ndim > 0 && ndim <= 26) << "Only support up to 26 dimensions";
+  ICHECK(ndim >= 0 && ndim <= 26) << "Only support up to 26 dimensions, but 
got " << ndim;
   return Layout("ABCDEFGHIJKLMNOPQRSTUVWXYZ").SubLayout(0, ndim);
 }
 
@@ -61,7 +61,7 @@ LayoutDecision InitialLayoutDecision(int ndim) {
   if (ndim == kUnknownNDim) {
     return LayoutDecision::InitUnknownDim();
   }
-  ICHECK(ndim >= 0 && ndim <= 26) << "Only support up to 26 dimensions";
+  ICHECK(ndim >= 0 && ndim <= 26) << "Only support up to 26 dimensions, but 
got " << ndim;
   return Layout("ABCDEFGHIJKLMNOPQRSTUVWXYZ").SubLayout(0, ndim);
 }
 
diff --git a/tests/python/relax/test_transform_convert_layout.py 
b/tests/python/relax/test_transform_convert_layout.py
index 78a1b166dd..5187ab30b7 100644
--- a/tests/python/relax/test_transform_convert_layout.py
+++ b/tests/python/relax/test_transform_convert_layout.py
@@ -708,6 +708,19 @@ def test_conv2d_transpose():
     verify(Input, Expected)
 
 
+def test_expand_dims_scalar():
+    @I.ir_module
+    class Input:
+        @R.function
+        def main() -> R.Tensor((1,), dtype="int64"):
+            with R.dataflow():
+                gv: R.Tensor((1,), dtype="int64") = R.expand_dims(R.const(0, 
"int64"), axis=[0])
+                R.output(gv)
+            return gv
+
+    verify(Input, Input)
+
+
 def test_conv2d_expand_dims():
     @I.ir_module
     class Input: