[API] Standardize MXNet NumPy creation functions

benchmark/python/sparse/cast_storage.py

@@ -41,7 +41,7 @@ def measure_cost(repeat, f, *args, **kwargs):
 
 def run_cast_storage_synthetic():
     def dense_to_sparse(m, n, density, ctx, repeat, stype):
-        set_default_context(ctx)
+        set_default_device(ctx)
         data_shape = (m, n)
         dns_data = rand_ndarray(data_shape, stype, density).tostype('default')
         dns_data.wait_to_read()

benchmark/python/sparse/dot.py

@@ -26,7 +26,7 @@
 import mxnet as mx
 import numpy as np
 import numpy.random as rnd
-from mxnet.test_utils import rand_ndarray, set_default_context, assert_almost_equal, get_bz2_data
+from mxnet.test_utils import rand_ndarray, set_default_device, assert_almost_equal, get_bz2_data
 from mxnet.base import check_call, _LIB
 from util import estimate_density
 
@@ -267,7 +267,7 @@ def test_dot_synthetic(data_dict):
     # Benchmark MXNet and Scipys dot operator
     def bench_dot(lhs_shape, rhs_shape, lhs_stype, rhs_stype,
                   lhs_den, rhs_den, trans_lhs, ctx, num_repeat=10, fw="mxnet", distribution="uniform"):
-        set_default_context(ctx)
+        set_default_device(ctx)
         assert fw == "mxnet" or fw == "scipy"
         # Set funcs
         dot_func_sparse = mx.nd.sparse.dot if fw == "mxnet" else sp.spmatrix.dot

benchmark/python/sparse/sparse_op.py

@@ -155,7 +155,7 @@ def measure_cost_backward_baseline(repeat, dot, transpose, lhs, rhs):
         return diff / repeat
 
     def bench_dot_forward(m, k, n, density, ctx, repeat):
-        set_default_context(ctx)
+        set_default_device(ctx)
         dns = mx.nd.random.uniform(shape=(k, n)).copyto(ctx)
         data_shape = (m, k)
         csr_data = rand_ndarray(data_shape, 'csr', density)
@@ -184,7 +184,7 @@ def bench_dot_forward(m, k, n, density, ctx, repeat):
                      ratio_baseline, costs_baseline[0], costs_baseline[1]))
 
     def bench_dot_backward(m, k, n, density, ctx, repeat):
-        set_default_context(ctx)
+        set_default_device(ctx)
         dns = mx.nd.random.uniform(shape=(m, n)).copyto(ctx)
         data_shape = (m, k)
         csr_data = rand_ndarray(data_shape, 'csr', density)

ci/docker/runtime_functions.sh

@@ -874,6 +874,8 @@ unittest_array_api_standardization() {
     export DMLC_LOG_STACK_TRACE_DEPTH=100
     python3 -m pytest --durations=50 --cov-report xml:tests_api.xml --verbose \
         array_api_tests/test_type_promotion.py::test_elementwise_function_two_arg_bool_type_promotion
+    python3 -m pytest --durations=50 --cov-report xml:tests_api.xml --verbose array_api_tests/test_creation_functions.py
+    python3 -m pytest --durations=50 --cov-report xml:tests_api.xml --verbose array_api_tests/test_indexing.py
     popd
 }
 

docs/python_docs/python/api/context/index.rst → docs/python_docs/python/api/device/index.rst

@@ -15,9 +15,9 @@
    specific language governing permissions and limitations
    under the License.
 
-mxnet.context
+mxnet.device
 =============
 
-.. automodule:: mxnet.context
+.. automodule:: mxnet.device
     :members:
     :autosummary:

docs/python_docs/python/api/index.rst

@@ -86,10 +86,10 @@ Gluon related modules
       Key value store interface of MXNet for parameter synchronization.
 
    .. card::
-      :title: mxnet.context
-      :link: mxnet/context/index.html
+      :title: mxnet.device
+      :link: mxnet/device/index.html
 
-      CPU and GPU context information.
+      CPU and GPU device information.
 
    .. card::
       :title: mxnet.profiler
@@ -116,10 +116,10 @@ Advanced modules
       API for querying MXNet enabled features.
 
    .. card::
-      :title: mxnet.context
-      :link: context/index.html
+      :title: mxnet.device
+      :link: device/index.html
 
-      MXNet array context for specifying in-memory storage device.
+      MXNet array device for specifying in-memory storage device.
 
    .. card::
       :title: mxnet.profiler

docs/python_docs/python/api/npx/index.rst

@@ -50,7 +50,7 @@ Devices
    cpu_pinned
    gpu
    gpu_memory_info
-   current_context
+   current_device
    num_gpus
 
 Nerual networks

docs/python_docs/python/tutorials/deploy/inference/image_classification_jetson.md

@@ -71,18 +71,18 @@ And we are done. You can test the installation now by importing mxnet from pytho
 
 ## Running a pre-trained ResNet-50 model on Jetson
 
-We are now ready to run a pre-trained model and run inference on a Jetson module. In this tutorial we are using ResNet-50 model trained on Imagenet dataset. We run the following classification script with either cpu/gpu context using python3.
+We are now ready to run a pre-trained model and run inference on a Jetson module. In this tutorial we are using ResNet-50 model trained on Imagenet dataset. We run the following classification script with either cpu/gpu device using python3.
 
 ```{.python .input}
 from mxnet import gluon
 import mxnet as mx
 
-# set context
+# set device
 gpus = mx.test_utils.list_gpus()
-ctx =  mx.gpu() if gpus else mx.cpu()
+device =  mx.gpu() if gpus else mx.cpu()
 
 # load pre-trained model
-net = gluon.model_zoo.vision.resnet50_v1(pretrained=True, ctx=ctx)
+net = gluon.model_zoo.vision.resnet50_v1(pretrained=True, device=device)
 net.hybridize(static_alloc=True, static_shape=True)
 
 # load labels
@@ -99,7 +99,7 @@ img = mx.image.color_normalize(img.astype(dtype='float32')/255,
                                std=mx.np.array([0.229, 0.224, 0.225])) # normalize
 img = img.transpose((2, 0, 1)) # channel first
 img = mx.np.expand_dims(img, axis=0) # batchify
-img = img.as_in_ctx(ctx)
+img = img.to_device(device)
 
 prob = mx.npx.softmax(net(img)) # predict and normalize output
 idx = mx.npx.topk(prob, k=5)[0] # get top 5 result

docs/python_docs/python/tutorials/extend/customop.md

@@ -101,7 +101,7 @@ class SigmoidProp(mx.operator.CustomOpProp):
         # return 3 lists representing inputs shapes, outputs shapes, and aux data shapes.
         return (data_shape,), (output_shape,), ()
 
-    def create_operator(self, ctx, in_shapes, in_dtypes):
+    def create_operator(self, device, in_shapes, in_dtypes):
         #  create and return the CustomOp class.
         return Sigmoid()
 ```
@@ -183,7 +183,7 @@ class DenseProp(mx.operator.CustomOpProp):
         # return 3 lists representing inputs shapes, outputs shapes, and aux data shapes.
         return (data_shape, weight_shape), (output_shape,), ()
 
-    def create_operator(self, ctx, in_shapes, in_dtypes):
+    def create_operator(self, device, in_shapes, in_dtypes):
         #  create and return the CustomOp class.
         return Dense(self._bias)
 ```
@@ -201,8 +201,8 @@ class DenseBlock(mx.gluon.Block):
         self.weight = gluon.Parameter('weight', shape=(channels, in_channels))
 
     def forward(self, x):
-        ctx = x.context
-        return mx.nd.Custom(x, self.weight.data(ctx), bias=self._bias, op_type='dense')
+        device = x.device
+        return mx.nd.Custom(x, self.weight.data(device), bias=self._bias, op_type='dense')
 ```
 
 ### Example usage

docs/python_docs/python/tutorials/getting-started/crash-course/1-nparray.md

@@ -195,7 +195,7 @@ a = np.array(a)
 (type(a), a)
 ```
 
-Additionally, you can move them to different GPU contexts. You will dive more
+Additionally, you can move them to different GPU devices. You will dive more
 into this later, but here is an example for now.
 
 ```{.python .input}

docs/python_docs/python/tutorials/getting-started/crash-course/2-create-nn.md

@@ -564,7 +564,7 @@ with warnings.catch_warnings():
     warnings.simplefilter("ignore")
     net_loaded = nn.SymbolBlock.imports("MLP_hybrid-symbol.json",
                                         ['data'], "MLP_hybrid-0000.params",
-                                        ctx=None)
+                                        device=None)
 ```
 
 ```{.python .input}

docs/python_docs/python/tutorials/getting-started/crash-course/4-components.md

@@ -35,7 +35,7 @@ import mxnet as mx
 from mxnet.gluon import nn
 npx.set_np()
 
-ctx = mx.cpu()
+device = mx.cpu()
 ```
 
 ## Initialization
@@ -103,7 +103,7 @@ To initialize your network using different built-in types, you have to use the
 from mxnet import init
 
 # Constant init initializes the weights to be a constant value for all the params
-net.initialize(init=init.Constant(3), ctx=ctx)
+net.initialize(init=init.Constant(3), device=device)
 print(net[0].weight.data()[0])
 ```
 
@@ -113,7 +113,7 @@ already initialized the weight but want to reinitialize the weight, set the
 `force_reinit` flag to `True`.
 
 ```{.python .input}
-net.initialize(init=init.Normal(sigma=0.2), force_reinit=True, ctx=ctx)
+net.initialize(init=init.Normal(sigma=0.2), force_reinit=True, device=device)
 print(net[0].weight.data()[0])
 ```
 

docs/python_docs/python/tutorials/getting-started/crash-course/5-datasets.md

@@ -215,7 +215,7 @@ The current data loading pipeline is the major bottleneck for many training task
 - `Transform.__call__()/forward()`
 - `Batchify`
 - (optional communicate through shared_mem)
-- `split_and_load(ctxs)`
+- `split_and_load(devices)`
 - training on GPUs
 
 Performance concerns include slow python dataset/transform functions, multithreading issues due to global interpreter lock, Python multiprocessing issues due to speed, and batchify issues due to poor memory management.

docs/python_docs/python/tutorials/getting-started/crash-course/6-train-nn.md

@@ -322,13 +322,13 @@ hybridize the model.
 
 ```{.python .input}
 # Create the model based on the blueprint provided and initialize the parameters
-ctx = mx.gpu()
+device = mx.gpu()
 
 initializer = mx.initializer.Xavier()
 
 model = LeafNetwork()
-model.initialize(initializer, ctx=ctx)
-model.summary(mx.np.random.uniform(size=(4, 3, 128, 128), ctx=ctx))
+model.initialize(initializer, device=device)
+model.summary(mx.np.random.uniform(size=(4, 3, 128, 128), device=device))
 model.hybridize()
 ```
 
@@ -368,7 +368,7 @@ def test(val_data):
     for batch in val_data:
         data = batch[0]
         labels = batch[1]
-        outputs = model(data.as_in_ctx(ctx))
+        outputs = model(data.to_device(device))
         acc.update([labels], [outputs])
 
     _, accuracy = acc.get()
@@ -396,8 +396,8 @@ for epoch in range(epochs):
         data = batch[0]
         label = batch[1]
         with mx.autograd.record():
-            outputs = model(data.as_in_ctx(ctx))
-            loss = loss_fn(outputs, label.as_in_ctx(ctx))
+            outputs = model(data.to_device(device))
+            loss = loss_fn(outputs, label.to_device(device))
         mx.autograd.backward(loss)
         trainer.step(batch_size)
         accuracy.update([label], [outputs])

docs/python_docs/python/tutorials/getting-started/crash-course/7-use-gpus.md

@@ -36,11 +36,11 @@ npx.num_gpus() #This command provides the number of GPUs MXNet can access
 
 ## Allocate data to a GPU
 
-MXNet's ndarray is very similar to NumPy's. One major difference is that MXNet's ndarray has a `context` attribute specifieing which device an array is on. By default, arrays are stored on `npx.cpu()`. To change it to the first GPU, you can use the following code, `npx.gpu()` or `npx.gpu(0)` to indicate the first GPU.
+MXNet's ndarray is very similar to NumPy's. One major difference is that MXNet's ndarray has a `device` attribute specifieing which device an array is on. By default, arrays are stored on `npx.cpu()`. To change it to the first GPU, you can use the following code, `npx.gpu()` or `npx.gpu(0)` to indicate the first GPU.
 
 ```{.python .input}
 gpu = npx.gpu() if npx.num_gpus() > 0 else npx.cpu()
-x = np.ones((3,4), ctx=gpu)
+x = np.ones((3,4), device=gpu)
 x
 ```
 
@@ -63,7 +63,7 @@ If you have multiple GPUs on your machine, MXNet can access each of them through
 To perform an operation on a particular GPU, you only need to guarantee that the input of an operation is already on that GPU. The output is allocated on the same GPU as well. Almost all operators in the `np` and `npx` module support running on a GPU.
 
 ```{.python .input}
-y = np.random.uniform(size=(3,4), ctx=gpu)
+y = np.random.uniform(size=(3,4), device=gpu)
 x + y
 ```
 
@@ -115,17 +115,17 @@ class LeafNetwork(nn.HybridBlock):
         return batch
 ```
 
-Load the saved parameters onto GPU 0 directly as shown below; additionally, you could use `net.collect_params().reset_ctx(gpu)` to change the device.
+Load the saved parameters onto GPU 0 directly as shown below; additionally, you could use `net.collect_params().reset_device(gpu)` to change the device.
 
 ```{.python .input}
 net = LeafNetwork()
-net.load_parameters('leaf_models.params', ctx=gpu)
+net.load_parameters('leaf_models.params', device=gpu)
 ```
 
 Use the following command to create input data on GPU 0. The forward function will then run on GPU 0.
 
 ```{.python .input}
-x = np.random.uniform(size=(1, 3, 128, 128), ctx=gpu)
+x = np.random.uniform(size=(1, 3, 128, 128), device=gpu)
 net(x)
 ```
 
@@ -201,7 +201,7 @@ devices = available_gpus[:num_gpus]
 print('Using {} GPUs'.format(len(devices)))
 
 # Diff 2: reinitialize the parameters and place them on multiple GPUs
-net.initialize(force_reinit=True, ctx=devices)
+net.initialize(force_reinit=True, device=devices)
 
 # Loss and trainer are the same as before
 loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()

docs/python_docs/python/tutorials/getting-started/gluon_from_experiment_to_deployment.md

@@ -91,10 +91,10 @@ lr_factor = 0.75
 # learning rate change at following epochs
 lr_epochs = [10, 20, 30]
 
-num_gpus = mx.context.num_gpus()
+num_gpus = mx.device.num_gpus()
 # you can replace num_workers with the number of cores on you device
 num_workers = 8
-ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]
+device = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]
 batch_size = per_device_batch_size * max(num_gpus, 1)
 ```
 
@@ -166,11 +166,11 @@ Before we go to training, one unique Gluon feature you should be aware of is hyb
 
 ```{.python .input}
 # load pre-trained resnet50_v2 from model zoo
-finetune_net = resnet50_v2(pretrained=True, ctx=ctx)
+finetune_net = resnet50_v2(pretrained=True, device=device)
 
 # change last softmax layer since number of classes are different
 finetune_net.output = nn.Dense(classes)
-finetune_net.output.initialize(init.Xavier(), ctx=ctx)
+finetune_net.output.initialize(init.Xavier(), device=device)
 # hybridize for better performance
 finetune_net.hybridize()
 
@@ -195,11 +195,11 @@ Now let's define the test metrics and start fine-tuning.
 
 
 ```{.python .input}
-def test(net, val_data, ctx):
+def test(net, val_data, device):
     metric = mx.gluon.metric.Accuracy()
     for i, (data, label) in enumerate(val_data):
-        data = gluon.utils.split_and_load(data, ctx_list=ctx, even_split=False)
-        label = gluon.utils.split_and_load(label, ctx_list=ctx, even_split=False)
+        data = gluon.utils.split_and_load(data, device, even_split=False)
+        label = gluon.utils.split_and_load(label, device, even_split=False)
         outputs = [net(x) for x in data]
         metric.update(label, outputs)
     return metric.get()
@@ -215,8 +215,8 @@ for epoch in range(1, epochs + 1):
 
     for i, (data, label) in enumerate(train_data):
         # get the images and labels
-        data = gluon.utils.split_and_load(data, ctx_list=ctx, even_split=False)
-        label = gluon.utils.split_and_load(label, ctx_list=ctx, even_split=False)
+        data = gluon.utils.split_and_load(data, device, even_split=False)
+        label = gluon.utils.split_and_load(label, device, even_split=False)
         with autograd.record():
             outputs = [finetune_net(x) for x in data]
             loss = [softmax_cross_entropy(yhat, y) for yhat, y in zip(outputs, label)]
@@ -229,12 +229,12 @@ for epoch in range(1, epochs + 1):
 
     _, train_acc = metric.get()
     train_loss /= num_batch
-    _, val_acc = test(finetune_net, val_data, ctx)
+    _, val_acc = test(finetune_net, val_data, device)
 
     print('[Epoch %d] Train-acc: %.3f, loss: %.3f | Val-acc: %.3f | learning-rate: %.3E | time: %.1f' %
           (epoch, train_acc, train_loss, val_acc, trainer.learning_rate, time.time() - tic))
 
-_, test_acc = test(finetune_net, test_data, ctx)
+_, test_acc = test(finetune_net, test_data, device)
 print('[Finished] Test-acc: %.3f' % (test_acc))
 ```