global numpy shape flag

ci/docker/runtime_functions.sh

@@ -1532,6 +1532,13 @@ nightly_estimator() {
     nosetests test_sentiment_rnn.py
 }
 
+nightly_global_numpy_flag() {
+    set -ex
+    export PYTHONPATH=/work/mxnet/python/
+    cd /work/mxnet/tests/nightly
+    nosetests test_global_numpy_shape.py
+}
+
 # For testing PRs
 deploy_docs() {
     set -ex

include/mxnet/c_api.h

@@ -1173,7 +1173,8 @@ MXNET_DLL int MXAutogradIsTraining(bool* curr);
 MXNET_DLL int MXIsNumpyShape(bool* curr);
 /*!
  * \brief set numpy compatibility switch
- * \param is_np_shape 1 when numpy shape semantics is on, 0 when off
+ * \param is_np_shape 1 when numpy shape semantics is thread local on,
+ *        2 when numpy shape semantics is global on and 0 when off
  * \param prev returns the previous status before this set
  * \return 0 when success, -1 when failure happens
  */

include/mxnet/imperative.h

@@ -35,6 +35,17 @@
 #include "./ndarray.h"
 
 namespace mxnet {
+  /*! \brief there are three numpy shape flags based on priority.
+   * GlobalOn
+   *   turn on numpy shape flag globally, it includes thread local.
+   *   The flag can be seen in any thread.
+   * ThreadLocalOn
+   *   only turn on thread local numpy shape flag, it cannot be seen
+   *   in other threads.
+   * Off
+   *   turn off numpy shape flag globally.
+   * */
+  enum NumpyShape{Off, ThreadLocalOn, GlobalOn};
 /*! \brief runtime functions for NDArray */
 class Imperative {
  public:
@@ -97,14 +108,30 @@ class Imperative {
       is_recording_ = is_recording;
       return old;
   }
-  /*! brief whether numpy compatibility is on. */
+  /*! \brief whether numpy compatibility is on. */
   bool is_np_shape() const {
-    return is_np_shape_;
+    if (is_np_shape_global_) {
+      return true;
+    }
+    return is_np_shape_thread_local_;
   }
-  /*! brief turn on or turn off numpy compatibility switch. */
-  bool set_is_np_shape(bool is_np_shape) {
-    bool old = is_np_shape_;
-    is_np_shape_ = is_np_shape;
+  /*! \brief specify numpy compatibility off, thread local on or global on. */
+  bool set_is_np_shape(int is_np_shape) {
+    NumpyShape flag = static_cast<NumpyShape>(is_np_shape);
+    bool old = this->is_np_shape();
+    switch (flag) {
+      case GlobalOn:
+        is_np_shape_global_ = true;
+        is_np_shape_thread_local_ = true;
+        break;
+      case ThreadLocalOn:
+        is_np_shape_thread_local_ = true;
+        break;
+      case Off:
+        is_np_shape_global_ = false;
+        is_np_shape_thread_local_ = false;
+        break;
+    }
     return old;
   }
   /*! \brief to record operator, return corresponding node. */
@@ -177,14 +204,15 @@ class Imperative {
   static thread_local bool is_recording_;
   // TOOD(junwu): Added numpy compatibility switch for backward compatibility.
   // Delete it in the next major release.
-  static thread_local bool is_np_shape_;
+  static thread_local bool is_np_shape_thread_local_;
 #else
   static MX_THREAD_LOCAL bool is_train_;
   static MX_THREAD_LOCAL bool is_recording_;
   // TOOD(junwu): Added numpy compatibility switch for backward compatibility.
   // Delete it in the next major release.
-  static MX_THREAD_LOCAL bool is_np_shape_;
+  static MX_THREAD_LOCAL bool is_np_shape_thread_local_;
 #endif
+  bool is_np_shape_global_{false};
   /*! \brief node count used for naming */
   std::atomic<uint64_t> node_count_{0};
   /*! \brief variable count used for naming */

src/c_api/c_api_ndarray.cc

@@ -284,7 +284,7 @@ int MXIsNumpyShape(bool* curr) {
 
 int MXSetIsNumpyShape(int is_np_shape, int* prev) {
   API_BEGIN();
-  *prev = Imperative::Get()->set_is_np_shape(static_cast<bool>(is_np_shape));
+  *prev = Imperative::Get()->set_is_np_shape(is_np_shape);
   API_END();
 }
 

src/imperative/imperative.cc

@@ -25,11 +25,11 @@ namespace mxnet {
 #if DMLC_CXX11_THREAD_LOCAL
 thread_local bool Imperative::is_train_ = false;
 thread_local bool Imperative::is_recording_ = false;
-thread_local bool Imperative::is_np_shape_ = false;
+thread_local bool Imperative::is_np_shape_thread_local_ = false;
 #else
 MX_THREAD_LOCAL bool Imperative::is_train_ = false;
 MX_THREAD_LOCAL bool Imperative::is_recording_ = false;
-MX_THREAD_LOCAL bool Imperative::is_np_shape_ = false;
+MX_THREAD_LOCAL bool Imperative::is_np_shape_thread_local_ = false;
 #endif
 
 Imperative* Imperative::Get() {

tests/nightly/JenkinsfileForBinaries

@@ -49,6 +49,15 @@ core_logic: {
         }
       }
     },
+    // only MXNet subproject use this flag, so test in nightly
+    'CPU: Global_numpy_flag': {
+       node(NODE_LINUX_CPU) {
+         ws('workspace/global-numpy-flag') {
+           utils.unpack_and_init('cpu', mx_lib)
+           utils.docker_run('ubuntu_nightly_cpu', 'nightly_global_numpy_flag', false)
+         }
+       }
+    },
     'GPU: USE_INT64_TENSOR_SIZE': {
       node(NODE_LINUX_GPU) {
         ws('workspace/build-gpu-int64') {

tests/nightly/test_global_numpy_shape.py

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import mxnet as mx
+import numpy as np
+import threading
+from mxnet.util import set_np_shape
+from mxnet.test_utils import assert_almost_equal
+
+# the test should be tested in isolated environment
+# without other tests running parallelly as the flag
+# effect every single thread
+def test_np_global_shape():
+    set_np_shape(2)
+    data = []
+
+    def f():
+        # scalar
+        data.append(mx.np.ones(shape=()))
+        # zero-dim
+        data.append(mx.np.ones(shape=(0, 1, 2)))
+
+    thread = threading.Thread(target=f)
+    thread.start()
+    thread.join()
+
+    assert_almost_equal(data[0].asnumpy(), np.ones(shape=()))
+    assert_almost_equal(data[1].asnumpy(), np.ones(shape=(0, 1, 2)))
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule()