For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

ParallelArrayOps

ParallelArrayOps

class max.nn.parallel.ParallelArrayOps(accelerator=None, max_workers=24)

source

Bases: object

Parallelized numpy array operations for performance-critical data processing.

Uses ThreadPoolExecutor to parallelize bulk copy operations that release the GIL, enabling true multi-threaded execution. Particularly effective for concatenating large arrays where memory bandwidth can be saturated across multiple cores.

Thread pool cleanup is handled automatically via __del__ and weakref.finalize, providing defense-in-depth for resource cleanup.

import numpy as np
from max.nn.parallel.numpy_parallel_ops import ParallelArrayOps

arr1 = np.ones((2, 4), dtype=np.float32)
arr2 = np.zeros((3, 4), dtype=np.float32)
arr3 = np.ones((1, 4), dtype=np.float32)

ops = ParallelArrayOps(max_workers=20)
result = ops.concatenate([arr1, arr2, arr3], axis=0)

Initialize parallel array operations with a thread pool.

Parameters:

• accelerator (Device | None) – The accelerator to allocate pinned memory on. If provided, the results of the concatenate operation will be allocated on a pinned buffer on the specified accelerator.
• max_workers (int) – Maximum number of worker threads. Default is 24, which works well for typical server CPUs. Consider setting to match your expected number of arrays (such as 20 for up to 20 concurrent copies).

concatenate()

concatenate(arrays, axis=0, min_chunk_size_mb=50.0)

source

Concatenate arrays in parallel along the specified axis.

Equivalent to np.concatenate but parallelized using thread pool. Automatically splits large arrays across multiple workers when there are fewer arrays than workers, maximizing memory bandwidth utilization. Uses intelligent heuristics to avoid splitting when overhead would exceed benefits.

This matches np.concatenate’s behavior of always returning a copy, even when arrays contains a single array.

Parameters:

• arrays (Sequence[ndarray[tuple[Any, ...], dtype[Any]]]) – List of numpy arrays to concatenate. Must have compatible shapes (same shape except along concatenation axis) and identical dtypes.
• axis (int) – Axis along which to concatenate. Negative values count from the end. Default is 0.
• min_chunk_size_mb (float) – Minimum chunk size in MB when splitting arrays. Default is 50 MB. Prevents creating too many small work items that would have excessive overhead. Arrays are only split if the resulting chunks would be at least this size.

Returns:

Concatenated array with the same dtype as the input arrays.

Raises:

• ValueError – If arrays is empty, shapes are incompatible, or dtypes differ.
• RuntimeError – If the thread pool has been shut down.

Return type:

Buffer

shutdown()

shutdown(wait=True)

source

Shutdown the thread pool and release resources.

Parameters:

wait (bool) – If True, wait for pending tasks to complete. If False, cancel pending tasks immediately.

Return type:

None