From 2b3aa852c5b5c9bb2ad49a8e18252978f350ef1e Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 12 Dec 2022 09:49:34 +0000
Subject: [PATCH] fix(imperative/dtype): use envvar to control behavior

GitOrigin-RevId: 8f11fa0c7e0ffec8f3dd9c7003fa9bebf8961e2b
---
 imperative/python/src/helper.cpp | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/imperative/python/src/helper.cpp b/imperative/python/src/helper.cpp
index ede45bd76..fc4848902 100644
--- a/imperative/python/src/helper.cpp
+++ b/imperative/python/src/helper.cpp
@@ -675,8 +675,17 @@ PyObject* dtype_mgb2np(mgb::DType dtype) {
         Py_XINCREF(Py_None);
         return Py_None;
     }
-    // NOTE: the following is additional
-    return reinterpret_cast<PyObject*>(descr.release());
+    static bool use_typeobj_as_dtype = MGB_GETENV("MGE_USE_TYPEOBJ_AS_DTYPE");
+    if (use_typeobj_as_dtype) {
+        if (dtype.has_param()) {
+            return reinterpret_cast<PyObject*>(descr.release());
+        }
+        PyObject* typeobj = reinterpret_cast<PyObject*>(descr->typeobj);
+        Py_XINCREF(typeobj);
+        return typeobj;
+    } else {
+        return reinterpret_cast<PyObject*>(descr.release());
+    }
 }
 
 mgb::DType dtype_np2mgb(PyObject* obj) {
-- 
GitLab