[incubator-mxnet] branch master updated: Fix default context handling in Gluon 2 (#19113)

[zhasheng]

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

zhasheng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/master by this push:
     new 131fa69  Fix default context handling in Gluon 2 (#19113)
131fa69 is described below

commit 131fa6931c6731d474bd7d68582f37947127ff65
Author: Leonard Lausen <lau...@amazon.com>
AuthorDate: Thu Sep 10 19:47:31 2020 -0700

    Fix default context handling in Gluon 2 (#19113)
---
 python/mxnet/gluon/block.py                    |  3 ++-
 tests/python/unittest/test_deferred_compute.py | 26 ++++++++++++++++++--------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py
index 3178c20..bac14ca 100644
--- a/python/mxnet/gluon/block.py
+++ b/python/mxnet/gluon/block.py
@@ -1411,7 +1411,8 @@ class HybridBlock(Block):
                 # HybridBlock is a child block of a HybridBlock that has been 
hybridized.
                 return super().__call__(x, *args)
 
-            return self._call_cached_op(x, *args)
+            with x.ctx:
+                return self._call_cached_op(x, *args)
 
     def forward(self, x, *args):
         """Defines the forward computation. Arguments can be either
diff --git a/tests/python/unittest/test_deferred_compute.py 
b/tests/python/unittest/test_deferred_compute.py
index 4b98470..1565fb1 100644
--- a/tests/python/unittest/test_deferred_compute.py
+++ b/tests/python/unittest/test_deferred_compute.py
@@ -379,7 +379,7 @@ def test_dc_numpy_indexing_error():
 ###############################################################################
 # Gluon
 ###############################################################################
-def _assert_dc_gluon(setup, net, setup_is_deterministic=True, numpy=True, 
autograd=True):
+def _assert_dc_gluon(setup, net, setup_is_deterministic=True, numpy=True, 
autograd=True, ctx=None):
     """Compare results of deferred compute and normal imperative mode.
 
     Parameters
@@ -397,9 +397,14 @@ def _assert_dc_gluon(setup, net, 
setup_is_deterministic=True, numpy=True, autogr
         Wrap in autograd
 
     """
+
     nd = mx.np if numpy else mx.nd
 
-    xs = setup(nd=nd)
+    if ctx is None:
+        ctx = mx.context.current_context()
+    with ctx:
+        xs = setup(nd=nd)
+
     ys = net(*xs)
     ys_np = [y.asnumpy() for y in ys]
 
@@ -439,13 +444,18 @@ def test_dc_hybridblock():
             assert x.shape[1] == 10  # due to in_units=10 above
             return self.dense(x) + self.weight.data(x.context)
 
-    net = MyBlock()
-    net.initialize()
-    _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=False)
-    with mx.util.np_array(True):
+    if mx.context.current_context() == mx.cpu(0):  # CPU tests
+        contexts = [mx.cpu(0), mx.cpu(1)]
+    else:  # Use default context, GPU tests
+        contexts = [mx.context.current_context()]
+    for ctx in contexts:
         net = MyBlock()
-        net.initialize()
-        _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=True)
+        net.initialize(ctx=contexts)
+        _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=False, ctx=ctx)
+        with mx.util.np_array(True):
+            net = MyBlock()
+            net.initialize(ctx=contexts)
+            _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=True, ctx=ctx)
 
 
 def test_dc_hybridblock_deferred_init_no_infer_shape_error():