Make MXIsNumpyShape return enum

apache · Oct 23, 2019 · 23b35bc · 23b35bc
1 parent 06b86da
commit 23b35bc
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 8 deletions.
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
@@ -1170,7 +1170,7 @@ MXNET_DLL int MXAutogradIsTraining(bool* curr);
  * \param curr returns the current status
  * \return 0 when success, -1 when failure happens
  */
-MXNET_DLL int MXIsNumpyShape(bool* curr);
+MXNET_DLL int MXIsNumpyShape(int* curr);
 /*!
  * \brief set numpy compatibility switch
  * \param is_np_shape 1 when numpy shape semantics is thread local on,

diff --git a/include/mxnet/imperative.h b/include/mxnet/imperative.h
@@ -108,12 +108,14 @@ class Imperative {
       is_recording_ = is_recording;
       return old;
   }
-  /*! \brief whether numpy compatibility is on. */
-  bool is_np_shape() const {
+  /*! \brief return current numpy compatibility status,
+   *  GlobalOn(2), ThreadLocalOn(1), Off(0).
+   * */
+  int is_np_shape() const {
     if (is_np_shape_global_) {
-      return true;
+      return 2;
     }
-    return is_np_shape_thread_local_;
+    return is_np_shape_thread_local_ ? 1 : 0;
   }
   /*! \brief specify numpy compatibility off, thread local on or global on. */
   bool set_is_np_shape(int is_np_shape) {

diff --git a/scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc b/scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc
@@ -2778,7 +2778,7 @@ JNIEXPORT jint JNICALL Java_org_apache_mxnet_LibInfo_mxDumpProfile
 JNIEXPORT jint JNICALL Java_org_apache_mxnet_LibInfo_mxIsNumpyShape
   (JNIEnv *env, jobject obj, jobject compatibleRef) {
   bool isNumpyShape;
-  int ret = MXIsNumpyShape(&isNumpyShape);
+  int ret = MXIsNumpyShape(static_cast<int*>(&isNumpyShape));
   SetIntField(env, compatibleRef, static_cast<int>(isNumpyShape));
   return ret;
 }

diff --git a/src/c_api/c_api_ndarray.cc b/src/c_api/c_api_ndarray.cc
@@ -276,7 +276,7 @@ int MXAutogradSetIsRecording(int is_recording, int* prev) {
   API_END();
 }
 
-int MXIsNumpyShape(bool* curr) {
+int MXIsNumpyShape(int* curr) {
   API_BEGIN();
   *curr = Imperative::Get()->is_np_shape();
   API_END();

diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
@@ -1732,7 +1732,8 @@ bool NDArray::Load(dmlc::Stream *strm) {
            " Please turn on np shape semantics in Python using `with np_shape(True)`"
            " or decorator `use_np_shape` to scope the code of loading the ndarray.";
   } else {
-    CHECK(!Imperative::Get()->is_np_shape())
+    // when the flag is global on, skip the check since it would be always global on.
+    CHECK(Imperative::Get()->is_np_shape() == GlobalOn || !Imperative::Get()->is_np_shape())
         << "ndarray was not saved in np shape semantics, but being loaded in np shape semantics."
            " Please turn off np shape semantics in Python using `with np_shape(False)`"
            " to scope the code of loading the ndarray.";