gpu_device=mx.gpu() # Change this to mx.cpu() in absence of GPUs.

def f():
    a = mx.nd.ones((100,100))
    b = mx.nd.ones((100,100))
    c = a + b
    print(c)
# in default mx.cpu() is used
f()
# change the default context to the first GPU
with mx.Context(gpu_device):
    f()

a = mx.nd.ones((100, 100), gpu_device)
a

a = mx.nd.ones((100,100), mx.cpu())
b = mx.nd.ones((100,100), gpu_device)
c = mx.nd.ones((100,100), gpu_device)
a.copyto(c)  # copy from CPU to GPU
d = b + c
e = b.as_in_context(c.context) + c  # same to above
{'d':d, 'e':e}
x = mx.nd.array([[1, 2, 3], [4, 5, 6]])
type(x)

x.shape
y = x + mx.nd.ones(x.shape)*3
print(y.asnumpy())
z = y.as_in_context(mx.gpu(0))
print(z)