Skip to main content
Log in

Python module

driver

CPU()

max.driver.CPU(id: int = -1) → Device

Creates a CPU device with the provided numa id.

CUDA()

max.driver.CUDA(id: int = -1) → Device

Creates a CUDA device with the provided id.

DLPackArray

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

Device

class max.driver.Device(_device: Device)

Device object. Limited to CUDA and CPU devices for now.

cpu()

classmethod cpu(id: int = -1) → Device

Creates a CPU device with the provided numa id.

cuda()

classmethod cuda(id: int = -1) → Device

Creates a CUDA device with the provided id.

id

property id*: int*

Returns device id.

is_host

property is_host

Returns whether the device is the CPU.

label

property label*: str*

Returns device label.

stats

property stats*: Mapping[str, Any]*

Returns utilization data for the device.

DeviceSpec

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

cpu()

static cpu(id: int = -1)

cuda()

static cuda(id: int = -1)

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(filename: PathLike, dtype: DType, shape: Sequence[int], mode='r+', offset=0)

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

The constructor argument semantics follow that of np.memmap.

read_only

property read_only*: bool*

Tensor

class max.driver.Tensor(shape: ~typing.Sequence[int], dtype: ~max.dtype.dtype.DType, device: ~max.driver.driver.Device = Device(_device=<max._driver.core.Device object>))

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.

contiguous()

contiguous() → Tensor

Creates a contiguous copy of the parent tensor.

copy()

copy(device: Device | None = None) → Tensor

Create a deep copy on an optionally given device.

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

device

property device*: Device*

Device on which tensor is resident.

dtype

property dtype*: DType*

DType of constituent elements in tensor.

element_size

property element_size*: int*

Return the size of the element type in bytes.

from_dlpack()

classmethod from_dlpack(arr: Any, *, 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()

classmethod from_numpy(arr: ndarray) → 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*: bool*

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. These situations include: * A tensor with negative steps.

is_host

property is_host*: bool*

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

item()

item() → Any

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*: int*

Returns the number of elements in this tensor.

Rank-0 tensors have 1 element by convention.

rank

property rank*: int*

Tensor rank.

scalar()

classmethod scalar(value: ~typing.Any, dtype: ~max.dtype.dtype.DType, device: ~max.driver.driver.Device = Device(_device=<max._driver.core.Device object>)) → Tensor

Create a scalar value of a given dtype and value.

shape

property shape*: Sequence[int]*

Shape of tensor.

to()

to(device: Device) → Tensor

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()

classmethod zeros(shape: ~typing.Sequence[int], dtype: ~max.dtype.dtype.DType, device: ~max.driver.driver.Device = Device(_device=<max._driver.core.Device object>)) → Tensor

Allocates an tensor with all elements initialized to zero.