Apache MXNet - Python API ndarray

Apache MXNet - Python API: NDArray

This chapter explores mxnet.ndarray (often imported as mx.nd), the core library in Apache MXNet for numerical computation. NDArrays are fundamental for handling data in MXNet.

Understanding mxnet.ndarray

MXNet’s NDArray provides a versatile N-dimensional array object, similar to NumPy arrays, but with key features tailored for deep learning and high-performance computing:

• Hardware Acceleration: NDArray operations can run on various hardware, including CPUs and NVIDIA GPUs, enabling significant speedups for large computations.
• Automatic Parallelization: MXNet’s execution engine automatically parallelizes operations on NDArrays across available hardware resources.
• Asynchronous Execution: Operations are often executed asynchronously, especially on GPUs, allowing the Python frontend to continue processing while computations run in the background.

Note: MXNet 2.0 introduced the mxnet.numpy API, which aims for greater compatibility with NumPy. While mxnet.ndarray (mx.nd) remains fundamental and is used throughout older codebases and internally, new projects might prefer mxnet.numpy for a more familiar NumPy-like experience. This tutorial focuses on the original mxnet.ndarray API.

Here’s how to create NDArrays from Python lists:

import mxnet as mx
from mxnet import nd

# Create a 1-D NDArray
x = nd.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
print(x)

Output:

<html>
 <body>
  <p>
   [ 1.  2.  3.  4.  5.  6.  7.  8.  9. 10.]
   <NDArray 10 @cpu(0)>
  </p>
 </body>
</html>

Example of a 2-D NDArray:

# Create a 2-D NDArray (3x10 matrix)
y = nd.array([[1,2,3,4,5,6,7,8,9,10],
              [11,12,13,14,15,16,17,18,19,20],
              [21,22,23,24,25,26,27,28,29,30]])
print(y)

Output:

<html>
 <body>
  <p>
   [[ 1.  2.  3.  4.  5.  6.  7.  8.  9. 10.]
    [11. 12. 13. 14. 15. 16. 17. 18. 19. 20.]
    [21. 22. 23. 24. 25. 26. 27. 28. 29. 30.]]
   <NDArray 3x10 @cpu(0)>
  </p>
 </body>
</html>

Let’s delve into the classes, functions, and parameters of the mxnet.ndarray API.

Core Classes

Class	Definition
NDArray(handle[, writable])	The primary array object representing a multi-dimensional, homogeneous array of fixed-size items. It can reside on CPU or GPU memory.
CachedOp(sym[, flags])	A handle for a cached operator, used internally for optimizing repeated execution of symbolic graphs with NDArray inputs. Typically not directly used by end-users.

Common Functions and Parameters

The mxnet.ndarray module provides a rich set of functions for numerical computation, similar to those found in NumPy. Many of these functions are also available as methods of the NDArray object itself (e.g., x.sum() is equivalent to nd.sum(x)).

Function & its Parameters	Definition
Activation(data, act_type=‘relu’, out=None, name=None)	Applies an activation function element-wise to the input. Supported act_type include ‘relu’, ‘sigmoid’, ‘tanh’, ‘softrelu’, ‘softsign’. This is commonly used in neural networks.
BatchNorm(data, gamma, beta, moving_mean, moving_var, eps=0.001, momentum=0.9, fix_gamma=False, use_global_stats=False, output_mean_var=False, axis=1, cudnn_off=False, min_calib_range=None, max_calib_range=None, out=None, name=None)	Performs batch normalization. Normalizes a data batch by its mean and variance, then applies a learned scale (gamma) and offset (beta). Crucial for stabilizing and accelerating training of deep neural networks.
BilinearSampler(data, grid, cudnn_off=False, out=None, name=None)	Applies bilinear sampling to an input feature map using a grid. This is a key component in Spatial Transformer Networks, allowing for learnable spatial transformations of feature maps. The grid specifies the sampling coordinates.
BlockGrad(data, out=None, name=None)	Stops gradient computation. During backpropagation, gradients will not flow through this operator to its inputs. Useful for freezing parts of a network or for specific adversarial training setups.
cast(data, dtype, out=None, name=None)	Casts all elements of the input NDArray to a new data type (dtype), e.g., ‘float32’, ‘float16’, ‘int32’.

Implementation Examples

Example using BilinearSampler to upsample (zoom in) a small feature map. Here, we define an affine transformation matrix for scaling.

import mxnet as mx
from mxnet import nd

# Input data (1 channel, 2x2 image)
data = nd.array([[[[1, 2],
                   [3, 4]]]], ctx=mx.cpu())

# Affine matrix for 2x scaling (0.5 in grid generation means zoom in by 2x)
# [s_x, 0,   0  ]
# [0,   s_y, 0  ]  (for target_shape relative to source)
# We want to map a 4x4 output grid to the 2x2 input
affine_matrix = nd.array([[0.5, 0, 0],
                          [0, 0.5, 0]], ctx=mx.cpu())

affine_matrix = affine_matrix.reshape((1, 6)) # Batch size 1, 2x3 matrix flattened

# Generate a grid for a 4x4 output, mapping back to the 'data' coordinates
grid = nd.GridGenerator(data=affine_matrix, transform_type='affine', target_shape=(4, 4))

output = nd.BilinearSampler(data, grid)
print(output)

Output (values will be interpolated):

<html>
 <body>
  <p>
   [[[[1.   1.25 1.75 2.  ]
     [1.5  1.75 2.25 2.5 ]
     [2.5  2.75 3.25 3.5 ]
     [3.   3.25 3.75 4.  ]]]]
   <NDArray 1x1x4x4 @cpu(0)>
  </p>
 </body>
</html>

Example of shifting data horizontally using BilinearSampler with a warp field. A positive horizontal shift in the warp grid means sampling from the left.

import mxnet as mx
from mxnet import nd

# Input data (1 channel, 4x4 image)
data = nd.array([[[[ 1,  2,  3,  4],
                   [ 5,  6,  7,  8],
                   [ 9, 10, 11, 12],
                   [13, 14, 15, 16]]]], ctx=mx.cpu())

# Create a warp grid. A value of 1 in the first channel of warp_matrix means a shift of 1 pixel.
# We want to shift the image content to the left by 1 pixel, so samples are taken from 1 pixel to the right.
# Grid values are normalized: -1 to 1. A shift of 1 pixel in a 4-pixel wide image is 1 * (2/4) = 0.5 in normalized coordinates.
# A positive value in the x-flow field shifts sampling to the right.
flow_x = nd.full((1,1,4,4), 0.5, ctx=mx.cpu()) # Shift samples by 1 pixel to the right
flow_y = nd.zeros((1,1,4,4), ctx=mx.cpu())
warp_matrix = nd.concat(flow_x, flow_y, dim=1) # Shape (1, 2, H, W)

grid = nd.GridGenerator(data=warp_matrix, transform_type='warp', target_shape=(4,4)) # target_shape is often ignored for warp
output = nd.BilinearSampler(data, grid)
print("Data shifted (content moves left, zeros appear on right):")
print(output)

Output (content shifted left, rightmost column becomes zeros due to sampling out of bounds, which BilinearSampler often handles as zero):

<html>
 <body>
  <p>
Data shifted (content moves left, zeros appear on right):
[[[[ 2.   3.   4.   0.  ]
   [ 6.   7.   8.   0.  ]
   [10.  11.  12.   0.  ]
   [14.  15.  16.   0.  ]]]]
   <NDArray 1x1x4x4 @cpu(0)>
  </p>
 </body>
</html>

Example using cast() function:

original_array = nd.array([300.0, 10.1, 15.4, -1.0, -2.0])
casted_array = nd.cast(original_array, dtype='uint8')
print(f"Original: {original_array}")
print(f"Casted to uint8: {casted_array}")

Output (values are clamped/wrapped to uint8 range [0, 255]):

<html>
 <body>
  <p>
Original: [300.   10.1  15.4  -1.   -2. ] <NDArray 5 @cpu(0)>
Casted to uint8: [255  10  15 255 254] <NDArray 5 @cpu(0)>
  </p>
 </body>
</html>

ndarray.contrib: Experimental Features

The mxnet.ndarray.contrib package hosts experimental APIs and features. These functions are provided for the community to try and give feedback. They might be promoted to the main mxnet.ndarray API, changed, or removed in future versions. Use with caution in production environments.

Selected Functions from ndarray.contrib

Function & its Parameters	Definition
rand_zipfian(true_classes, num_sampled, range_max, out=None, name=None)	Draws random samples from an approximate Zipfian distribution. Often used in candidate sampling for large-scale classification tasks (e.g., in natural language processing).
foreach(body, data, init_states, name=None)	Executes a for-loop-like computation over slices of NDArrays. The body is a Python function defining the computation for each iteration. This is a way to express certain sequential operations that can still be managed by the MXNet engine.
while_loop(cond, func, loop_vars, max_iterations=None, name=None)	Executes a while-loop-like computation. cond is a Python function determining the loop condition, func defines the computation within the loop, and loop_vars are the variables iterated upon.
cond(pred, then_func, else_func, name=None)	Executes an if-then-else branch. pred is an NDArray scalar condition; then_func is executed if true, else_func if false. Both then_func and else_func are Python functions returning lists of NDArrays.
isinf(data, out=None, name=None)	Performs an element-wise check to determine if NDArray elements are infinite. Returns a boolean NDArray.
getnnz(data, axis=None, out=None, name=None)	Counts the number of non-zero stored values in a sparse tensor (specifically CSR format on CPU). Includes explicit zeros if they are stored.
requantize(data, min_range, max_range, out_min_range=None, out_max_range=None, out_type=‘int8’, out=None, name=None)	Requantizes data from int32 (with its min/max range) to a lower precision type like int8, using new min/max thresholds. Used in model quantization for inference optimization.

Implementation Examples (contrib)

Example using rand_zipfian:

from mxnet.contrib import nd as contrib_nd

true_classes = mx.nd.array([2, 7]) # Batch of true class labels
num_sampled = 3
range_max = 10 # Sample from [0, 9]

samples, expected_count_true, expected_count_sample = contrib_nd.rand_zipfian(
    true_classes, num_sampled, range_max
)
print(f"Samples:\n{samples}")
print(f"Expected count for true classes:\n{expected_count_true}")
print(f"Expected count for sampled classes:\n{expected_count_sample}")

Output (actual samples will vary):

<html>
 <body>
  <p>
Samples:
[[0 1 4]
 [3 0 1]]
<NDArray 2x3 @cpu(0)>
Expected count for true classes:
[0.411263 पैसा  0.23979962]
<NDArray 2 @cpu(0)>
Expected count for sampled classes:
[[0.73500365 0.5108256  0.3413174 ]
 [0.37796447 0.73500365 0.5108256 ]]
<NDArray 2x3 @cpu(0)>
  </p>
 </body>
</html>

Example using while_loop to compute a Fibonacci-like sequence up to a limit:

from mxnet.contrib import nd as contrib_nd

# Loop condition: current value < 20
cond = lambda i, prev, curr: curr < 20

# Loop function: update prev, curr; output curr
func = lambda i, prev, curr: ([curr], [i + 1, curr, prev + curr])

# Initial loop variables: (iteration_count, previous_value, current_value)
loop_vars = (mx.nd.array([0]), mx.nd.array([0]), mx.nd.array([1]))

outputs, states = contrib_nd.while_loop(cond, func, loop_vars, max_iterations=10)

print(f"Outputs from each iteration (current_value):\n{outputs[0]}")
print(f"Final states (i, prev, curr):\n{states}")

Output:

<html>
 <body>
  <p>
Outputs from each iteration (current_value):
[[ 1.]
 [ 1.]
 [ 2.]
 [ 3.]
 [ 5.]
 [ 8.]
 [13.]]
<NDArray 7x1 @cpu(0)>
Final states (i, prev, curr):
[<NDArray 1 @cpu(0)>, <NDArray 1 @cpu(0)>, <NDArray 1 @cpu(0)>]
  </p>
 </body>
</html>

ndarray.image: Image Manipulation Utilities

The mxnet.ndarray.image module provides functions specifically for image processing and augmentation tasks performed on NDArrays. These are often used in data preprocessing pipelines for computer vision models.

Selected Functions from ndarray.image

Function & its Parameters	Definition
adjust_lighting(data, alpha, out=None, name=None)	Adjusts image lighting using PCA-based noise, a technique from AlexNet. alpha is a vector of PCA component multipliers.
crop(data, x, y, width, height, out=None, name=None)	Crops an image NDArray. Input shape is typically (H x W x C) or (N x H x W x C). x, y are top-left coordinates for the crop.
normalize(data, mean=(0,0,0), std=(1,1,1), out=None, name=None)	Normalizes an image tensor of shape (C x H x W) or (N x C x H x W) by subtracting mean and dividing by std per channel.
random_crop(data, size, interp=2, out=None, name=None)	Randomly crops an image to a specified size (width, height). If source is smaller, it might upsample. interp controls interpolation method.
random_lighting(data, alpha_std=0.05, out=None, name=None)	Adds random PCA-based lighting noise to an image, following AlexNet style. alpha_std controls the standard deviation of the noise factors.
random_resized_crop(data, size, area=(0.08, 1.0), ratio=(0.75, 1.33), interp=2, out=None, name=None)	Crops a random region of the image and resizes it to size. The crop area and aspect ratio are randomized within specified area and ratio ranges. Common augmentation for ImageNet training.
resize(data, size, keep_ratio=False, interp=2, out=None, name=None)	Resizes an image to a target size (width, height). keep_ratio preserves aspect ratio. interp specifies interpolation method.
to_tensor(data, out=None, name=None)	Converts an image NDArray of shape (H x W x C) or (N x H x W x C) with pixel values in [0, 255] to a tensor of shape (C x H x W) or (N x C x H x W) with values in [0, 1] by dividing by 255.

Implementation Examples (image)

Example using to_tensor and normalize:

import mxnet as mx
from mxnet import nd
import numpy as np # For creating initial numpy array

# Create a dummy image (Height=2, Width=2, Channels=3) with uint8 values [0, 255]
dummy_image_np = np.array([[[10, 20, 30], [40, 50, 60]],
                           [[70, 80, 90], [100, 110, 120]]], dtype=np.uint8)
img_nd = mx.nd.array(dummy_image_np)
print(f"Original image (HWC):\n{img_nd}")

# Convert to tensor (CHW, float32, range [0,1])
img_tensor = nd.image.to_tensor(img_nd)
print(f"Image as tensor (CHW):\n{img_tensor}")

# Normalize the tensor
mean_vals = nd.array([0.485, 0.456, 0.406]) # Typical ImageNet means
std_vals = nd.array([0.229, 0.224, 0.225])  # Typical ImageNet stds
normalized_tensor = nd.image.normalize(img_tensor, mean=mean_vals, std=std_vals)
print(f"Normalized tensor:\n{normalized_tensor}")

Output (values will be approximate due to float precision):

<html>
 <body>
  <p>
Original image (HWC):
[[[ 10  20  30]
  [ 40  50  60]]

 [[ 70  80  90]
  [100 110 120]]]
<NDArray 2x2x3 @cpu(0)>
Image as tensor (CHW):
[[[0.03921569 0.15686275]
  [0.27450982 0.39215687]]

 [[0.07843138 0.19607843]
  [0.3137255  0.43137255]]

 [[0.11764706 0.23529412]
  [0.3529412  0.47058824]]]
<NDArray 3x2x2 @cpu(0)>
Normalized tensor:
[[[-1.9469869  -1.4330131 ]
  [-0.91903937 -0.40506548]]

 [[-1.6856009  -1.1579224 ]
  [-0.6302439  -0.10256535]]

 [[-1.2815042  -0.76430035]
  [-0.24709654  0.2701073 ]]]
<NDArray 3x2x2 @cpu(0)>
  </p>
 </body>
</html>

ndarray.random: Random Number Generation

The mxnet.ndarray.random module provides functions for generating NDArrays with random numbers drawn from various statistical distributions. This is essential for initializing model parameters, creating random data for testing, or implementing stochastic processes.

Selected Functions from ndarray.random

Function & its Parameters	Definition
uniform(low=0, high=1, shape=None, dtype=None, ctx=None, out=None)	Generates samples from a uniform distribution over [low, high).
normal(loc=0, scale=1, shape=None, dtype=None, ctx=None, out=None)	Generates samples from a normal (Gaussian) distribution with mean loc and standard deviation scale.
randn(*shape, **kwargs)	Generates samples from a standard normal distribution (mean 0, variance 1). shape is passed as separate arguments.
exponential(scale=1, shape=None, dtype=None, ctx=None, out=None)	Generates samples from an exponential distribution with rate 1/scale (lambda).
gamma(alpha, beta, shape=None, dtype=None, ctx=None, out=None)	Generates samples from a gamma distribution with shape alpha and scale beta.
multinomial(data, shape=1, get_prob=False, out=None, dtype=‘int32’)	Generates samples from multinomial distributions. data contains probabilities for each outcome.
negative_binomial(k, p, shape=None, dtype=None, ctx=None, out=None)	Generates samples from a negative binomial distribution (number of successes k before r failures, with success probability p). The interpretation of k, p can vary; check MXNet docs for specifics.
generalized_negative_binomial(mu, alpha, shape=None, dtype=None, ctx=None, out=None)	Generates samples from a generalized negative binomial distribution parameterized by mean mu and dispersion alpha.
shuffle(data, out=None, name=None)	Shuffles the elements of an NDArray along its first axis.
randint(low, high, shape=None, dtype=‘int32’, ctx=None, out=None)	Generates random integers from a discrete uniform distribution in the range [low, high).
*_like functions (e.g., uniform_like, normal_like)	These functions (e.g., normal_like(data, loc=0, scale=1)) generate random samples with the same shape and context as the input data NDArray, but with values drawn from the specified distribution.

Implementation Examples (random)

Drawing samples from uniform and normal distributions:

from mxnet.ndarray import random as random_nd # Alias for clarity

# Uniform distribution
uniform_samples = random_nd.uniform(low=-1, high=1, shape=(2, 3))
print(f"Uniform samples [-1, 1]:\n{uniform_samples}")

# Normal distribution
normal_samples = random_nd.normal(loc=0, scale=0.01, shape=(2,2))
print(f"Normal samples (mean 0, std 0.01):\n{normal_samples}")

# randint
int_samples = random_nd.randint(low=0, high=10, shape=(5,), ctx=mx.cpu())
print(f"Random integers [0, 9]:\n{int_samples}")

Output (actual random values will vary):

<html>
 <body>
  <p>
Uniform samples [-1, 1]:
[[-0.5206176   0.2503767  -0.95505637]
 [-0.7613349   0.04832268  0.9208205 ]]
<NDArray 2x3 @cpu(0)>
Normal samples (mean 0, std 0.01):
[[ 0.00133301 -0.00306308]
 [-0.00209531  0.00401994]]
<NDArray 2x2 @cpu(0)>
Random integers [0, 9]:
[8 3 0 9 2]
<NDArray 5 @cpu(0)>
  </p>
 </body>
</html>

ndarray.utils: Utility Functions

The mxnet.ndarray.utils module (often available directly under mx.nd for common functions like zeros) provides utility functions for creating, saving, and loading NDArrays.

Selected Functions from ndarray.utils

Function & its Parameters	Definition
zeros(shape, ctx=None, dtype=None, stype=‘default’)	Returns a new array of a given shape and dtype, filled with zeros. stype can specify sparse types like ‘csr’ or ‘row_sparse’.
empty(shape, ctx=None, dtype=None, stype=‘default’)	Returns a new array of a given shape and dtype, without initializing its entries. The content will be arbitrary memory values. Faster than zeros if you plan to fill it immediately.
array(source_array, ctx=None, dtype=None)	Creates an NDArray from an object exposing the array interface (e.g., Python list, NumPy array).
load(fname)	Loads a list or dictionary of NDArrays from a file. The file is typically saved using mx.nd.save().
load_frombuffer(buf)	Loads an NDArray dictionary or list from a binary string buffer (e.g., one read from a file or network).
save(fname, data)	Saves a list of NDArrays or a dictionary mapping strings to NDArrays into a binary file. data can be a single NDArray, a list of NDArrays, or a dict of str -> NDArray.

Implementation Examples (utils)

Creating arrays with zeros and using save/load:

# zeros is often accessed via mx.nd.zeros
x = mx.nd.zeros((2, 3), ctx=mx.cpu(), dtype='float32')
print(f"Zeros array:\n{x}")

# Save and load NDArrays
a = mx.nd.array([1, 2, 3])
b = mx.nd.array([[4, 5], [6, 7]])

file_path = 'my_arrays.params' # Common extension for saved NDArrays
# Save as a list
mx.nd.save(file_path, [a, b])
loaded_list = mx.nd.load(file_path)
print(f"Loaded list: {loaded_list}")

# Save as a dictionary
mx.nd.save(file_path, {'first': a, 'second': b})
loaded_dict = mx.nd.load(file_path)
print(f"Loaded dict: {loaded_dict}")

Output:

<html>
 <body>
  <p>
Zeros array:
[[0. 0. 0.]
 [0. 0. 0.]]
<NDArray 2x3 @cpu(0)>
Loaded list: [<NDArray 3 @cpu(0)>, <NDArray 2x2 @cpu(0)>]
Loaded dict: {'first': <NDArray 3 @cpu(0)>, 'second': <NDArray 2x2 @cpu(0)>}
  </p>
 </body>
</html>