Python module

driver

Exposes APIs for interacting with hardware, such as allocating tensors on a GPU and moving tensors between the CPU and GPU. It provides interfaces for memory management, device properties, and hardware monitoring. Through these APIs, you can control data placement, track resource utilization, and configure device settings for optimal performance.

`Accelerator`

class max.driver.Accelerator(self, id: int = -1)

Creates an accelerator device with the specified ID.

Provides access to GPU or other hardware accelerators in the system.

from max import driver
device = driver.Accelerator()
# Or specify GPU id
device = driver.Accelerator(id=0)  # First GPU
device = driver.Accelerator(id=1)  # Second GPU
# Get device id
device_id = device.id
from max import driver
device = driver.Accelerator()
# Or specify GPU id
device = driver.Accelerator(id=0)  # First GPU
device = driver.Accelerator(id=1)  # Second GPU
# Get device id
device_id = device.id

`CPU`

class max.driver.CPU(self, id: int = -1)

Creates a CPU device.

from max import driver
# Create default CPU device
device = driver.CPU()
# Device id is always 0 for CPU devices
device_id = device.id
from max import driver
# Create default CPU device
device = driver.CPU()
# Device id is always 0 for CPU devices
device_id = device.id

`DLPackArray`

class max.driver.DLPackArray(*args, **kwargs)

`Device`

class max.driver.Device

`api`

property api

Returns the API used to program the device.

Possible values are:

“cpu” for host devices
“cuda” for NVIDIA GPUs
“hip” for AMD GPUs

from max import driver

device = driver.CPU()
device.api
from max import driver

device = driver.CPU()
device.api

`can_access`

can_access

Checks if this device can directly access memory of another device.

Parameters:

other (Device) – The other device to check peer access against.
Returns:

True if peer access is possible, False otherwise.
Return type:

bool

Example

from max import driver

gpu0 = driver.Accelerator(id=0)
gpu1 = driver.Accelerator(id=1)

if gpu0.can_access(gpu1):
    print("GPU0 can directly access GPU1 memory.")
from max import driver

gpu0 = driver.Accelerator(id=0)
gpu1 = driver.Accelerator(id=1)

if gpu0.can_access(gpu1):
    print("GPU0 can directly access GPU1 memory.")

`cpu`

cpu = <nanobind.nb_func object>

`id`

property id

Returns a zero-based device id. For a CPU device this is always 0. For GPU accelerators this is the id of the device relative to this host. Along with the label, an id can uniquely identify a device, e.g. “gpu:0”, “gpu:1”.

from max import driver

device = driver.Accelerator()
device.id
from max import driver

device = driver.Accelerator()
device.id

`is_compatible`

property is_compatible

Returns whether this device is compatible with MAX.

`is_host`

property is_host

Whether this device is the CPU (host) device.

from max import driver

device = driver.CPU()
device.is_host
from max import driver

device = driver.CPU()
device.is_host

`label`

property label

Returns device label.

Possible values are:

“cpu” for host devices
“gpu” for accelerators

from max import driver

device = driver.CPU()
device.label
from max import driver

device = driver.CPU()
device.label

`stats`

property stats

Returns utilization data for the device.

from max import driver

device = driver.CPU()
device.stats
from max import driver

device = driver.CPU()
device.stats

`synchronize`

synchronize

Ensures all operations on this device complete before returning.

Raises:

ValueError – If any enqueue’d operations had an internal error.

`DeviceSpec`

class max.driver.DeviceSpec(id: int, device_type: Literal['cpu', 'gpu'] = 'cpu')

Specification for a device, containing its ID and type.

This class provides a way to specify device parameters like ID and type (CPU/GPU) for creating Device instances.

`accelerator()`

static accelerator(id: int = 0)

Creates an accelerator (GPU) device specification.

`cpu()`

static cpu(id: int = -1)

Creates a CPU device specification.

`device_type`

device_type*: Literal['cpu', 'gpu']* = 'cpu'

Type of specified device.

`id`

id*: int*

Provided id for this device.

`MemMapTensor`

class max.driver.MemMapTensor(self, shape: collections.abc.Sequence[int], dtype: max._core.dtype.DType, device: max._core.driver.Device | None = None)

`MemMapTensor`

class max.driver.MemMapTensor(self, shape: ndarray[writable=False], device: max._core.driver.Device)

`MemMapTensor`

class max.driver.MemMapTensor(self, other: max._core.driver.Tensor)

Create a memory-mapped tensor from a binary file on disk.

The constructor argument semantics follow that of np.memmap.

Overloaded function.

__init__(self, shape: collections.abc.Sequence[int], dtype: max._core.dtype.DType, device: max._core.driver.Device | None = None) -> None
__init__(self, shape: ndarray[writable=False], device: max._core.driver.Device) -> None
__init__(self, other: max._core.driver.Tensor) -> None

Moves the internals from an existing Tensor object into a new Tensor object.

Primarily used for initializing subclasses with existing Tensors.

`read_only`

property read_only*: bool*

`view()`

view(dtype: DType, shape: Sequence[int] | None = None) → MemMapTensor

Creates a view of the memory-mapped tensor with new type and shape.

Preserves the original memory mapping properties (read-only status and file reference) while reinterpreting the bytes. Shares underlying storage – modifications affect all views of the same memory region.

Parameters:
- dtype – New data type for interpreting the bytes. Size must divide evenly into the original tensor’s last dimension byte size.
- shape – Optional new shape. Inferred from dtype size if not provided. Product of dimensions must match original element count adjusted for dtype size.
Returns:

New MemMapTensor instance with specified dtype/shape, sharing the original memory-mapped file properties.
Raises:
- ValueError – For invalid shape/dtype combinations or boundary overflows.
- TypeError – For undefined byte interpretations.

`Tensor`

class max.driver.Tensor(self, shape: collections.abc.Sequence[int], dtype: max._core.dtype.DType, device: max._core.driver.Device | None = None)

`Tensor`

class max.driver.Tensor(self, shape: ndarray[writable=False], device: max._core.driver.Device)

`Tensor`

class max.driver.Tensor(self, other: max._core.driver.Tensor)

Device-resident tensor representation. Allocates memory onto a given device with the provided shape and dtype. Tensors can be sliced to provide strided views of the underlying memory, but any tensors input into model execution must be contiguous. Does not currently support setting items across multiple indices, but does support numpy-style slicing.

Parameters:
- dtype – DType of tensor
- shape – Tuple of positive, non-zero integers denoting the tensor shape.
- device – Device to allocate tensor onto.

Overloaded function.

__init__(self, shape: collections.abc.Sequence[int], dtype: max._core.dtype.DType, device: max._core.driver.Device | None = None) -> None
__init__(self, shape: ndarray[writable=False], device: max._core.driver.Device) -> None
__init__(self, other: max._core.driver.Tensor) -> None

Moves the internals from an existing Tensor object into a new Tensor object.

Primarily used for initializing subclasses with existing Tensors.

`contiguous()`

contiguous() → Tensor

Creates a contiguous copy of the parent tensor.

`copy`

copy

Create a deep copy on an optionally given device.

If a device is None (default), a copy is created on the same device.

from max import driver
from max.dtype import DType

cpu_tensor = driver.Tensor([2, 3], dtype=DType.bfloat16, device=driver.CPU())

cpu_copy = cpu_tensor.copy()
from max import driver
from max.dtype import DType

cpu_tensor = driver.Tensor([2, 3], dtype=DType.bfloat16, device=driver.CPU())

cpu_copy = cpu_tensor.copy()

`device`

property device

Device on which tensor is resident.

`dtype`

property dtype

DType of constituent elements in tensor.

`element_size`

property element_size

Return the size of the element type in bytes.

`from_dlpack()`

from_dlpack(*, copy: bool | None = None) → Tensor

Create a tensor from an object implementing the dlpack protocol.

This usually does not result in a copy, and the producer of the object retains ownership of the underlying memory.

`from_numpy()`

from_numpy() → Tensor

Creates a tensor from a provided numpy array on the host device.

The underlying data is not copied unless the array is noncontiguous. If it is, a contiguous copy will be returned.

`is_contiguous`

property is_contiguous

Whether or not tensor is contiguously allocated in memory. Returns false if the tensor is a non-contiguous slice.

Currently, we consider certain situations that are contiguous as non-contiguous for the purposes of our engine, such as when a tensor has negative steps.

`is_host`

property is_host

Whether or not tensor is host-resident. Returns false for GPU tensors, true for CPU tensors.

from max import driver
from max.dtype import DType

cpu_tensor = driver.Tensor([2, 3], dtype=DType.bfloat16, device=driver.CPU())

print(cpu_tensor.is_host)
from max import driver
from max.dtype import DType

cpu_tensor = driver.Tensor([2, 3], dtype=DType.bfloat16, device=driver.CPU())

print(cpu_tensor.is_host)

`item`

item

Returns the scalar value at a given location. Currently implemented only for zero-rank tensors. The return type is converted to a Python built-in type.

`num_elements`

property num_elements

Returns the number of elements in this tensor.

Rank-0 tensors have 1 element by convention.

`rank`

property rank

Tensor rank.

`scalar`

scalar = <nanobind.nb_func object>

`shape`

property shape

Shape of tensor.

`to`

to

Return a tensor that’s guaranteed to be on the given device.

The tensor is only copied if the input device is different from the device upon which the tensor is already resident.

`to_numpy()`

to_numpy() → ndarray

Converts the tensor to a numpy array.

If the tensor is not on the host, an exception is raised.

`view()`

view(dtype: DType, shape: Sequence[int] | None = None) → Tensor

Return a new tensor with the given type and shape that shares the underlying memory.

If the shape is not given, it will be deduced if possible, or a ValueError is raised.

`zeros`

zeros = <nanobind.nb_func object>

`accelerator_api()`

max.driver.accelerator_api() → str

Returns the API used to program the accelerator.

`devices_exist()`

max.driver.devices_exist(devices: list[max.driver.driver.DeviceSpec]) → bool

Identify if devices exist.

`load_devices()`

max.driver.load_devices(device_specs: list[max.driver.driver.DeviceSpec]) → list[max._core.driver.Device]

Initialize and return a list of devices, given a list of device specs.

`scan_available_devices()`

max.driver.scan_available_devices() → list[max.driver.driver.DeviceSpec]

Returns all accelerators if available, else return cpu.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

Accelerator​

CPU​

DLPackArray​

Device​

api​

can_access​

Example​

cpu​

id​

is_compatible​

is_host​

label​

stats​

synchronize​

DeviceSpec​

accelerator()​

cpu()​

device_type​

id​

MemMapTensor​

MemMapTensor​

MemMapTensor​

read_only​

view()​

Tensor​

Tensor​

Tensor​

contiguous()​

copy​

device​

dtype​

element_size​

from_dlpack()​

from_numpy()​

is_contiguous​

is_host​

item​

num_elements​

rank​

scalar​

shape​

to​

to_numpy()​

view()​

zeros​

accelerator_api()​

devices_exist()​

load_devices()​

scan_available_devices()​

`Accelerator`

`CPU`

`DLPackArray`

`Device`

`api`

`can_access`

Example

`cpu`

`id`

`is_compatible`

`is_host`

`label`

`stats`

`synchronize`

`DeviceSpec`

`accelerator()`

`cpu()`

`device_type`

`id`

`MemMapTensor`

`MemMapTensor`

`MemMapTensor`

`read_only`

`view()`

`Tensor`

`Tensor`

`Tensor`

`contiguous()`

`copy`

`device`

`dtype`

`element_size`

`from_dlpack()`

`from_numpy()`

`is_contiguous`

`is_host`

`item`

`num_elements`

`rank`

`scalar`

`shape`

`to`

`to_numpy()`

`view()`

`zeros`

`accelerator_api()`

`devices_exist()`

`load_devices()`

`scan_available_devices()`