- Fix unit test

Author: mbs-octoml

include/tvm/tir/stmt.h

@@ -599,9 +599,9 @@ class AllocateConstNode : public StmtNode {
   /*! \brief The optional data associated to the constant.
    */
   Optional<runtime::NDArray> data;
-  /*! \brief If the PrimFunc containing the Stmt is added to IRModule,
-       this is an optional index to indicate the index within
-       "Constants" attribute, that is a Array<NDArray> of IRModule.
+  /*!
+   * \brief If the PrimFunc containing the Stmt is added to IRModule, this is an optional index
+   * to indicate the index within "constants" attribute, that is a Array<NDArray> of IRModule.
    */
   Optional<Integer> irmod_storage_idx;
   /*! \brief The type of the buffer. */

python/tvm/tir/stmt.py

@@ -358,7 +358,7 @@ class AllocateConst(Stmt):
     data_or_idx : Union[NDArray, int]
         If an NDArray, this is the const data associated with the
         constant.  If an integer, this is the index into the
-        "Constants" attribute of the `IRModule` that contains the
+        "constants" attribute of the `IRModule` that contains the
         `AllocateConst`.
 
     body : Stmt

tests/python/unittest/test_custom_datatypes.py

@@ -21,6 +21,7 @@
 import pytest
 import tvm
 import tvm.topi.testing
+import tvm.testing
 from tvm import relay
 from tvm.relay.testing.layers import batch_norm_infer
 from tvm.target.datatype import (
@@ -560,4 +561,4 @@ def test_posites2():
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    tvm.testing.main()

tests/python/unittest/test_tir_transform_extract_constants.py

@@ -18,6 +18,7 @@
 import tvm
 from tvm import tir
 from tvm.script import tir as T
+import tvm.testing
 
 
 @tvm.script.ir_module
@@ -49,7 +50,7 @@ def constant3(a: T.handle) -> None:
 
 def test_const_extraction():
     mod = tvm.tir.transform.ExtractPrimFuncConstants()(Module4)
-    constants = mod.attrs["Constants"]
+    constants = mod.attrs["constants"]
     assert len(constants) == 2
 
     def _visit(stmt):
@@ -63,4 +64,4 @@ def _visit(stmt):
 
 
 if __name__ == "__main__":
-    test_const_extraction()
+    tvm.testing.main()