Python class

DeviceRef

`DeviceRef`

class max.graph.DeviceRef(device_type, id=0)

source

Bases: object

A symbolic device representation.

DeviceRef type representation consists of a DeviceKind and an id. This is a direct representation of the device attribute in MLIR.

The following example demonstrates how to create and use device references:

from max.graph import DeviceRef
gpu_device = DeviceRef.GPU()
print(gpu_device)  # Outputs: gpu:0
# Create a CPU device with specific id
cpu_device = DeviceRef.CPU(id=1)
print(cpu_device)  # Outputs: cpu:1

Parameters:

device_type (DeviceKind)
id (int)

`CPU()`

static CPU(id=0)

source

Creates a CPU device reference.

Parameters:: id (int)
Return type:: DeviceRef

`GPU()`

static GPU(id=0)

source

Creates a GPU device reference.

Parameters:: id (int)
Return type:: DeviceRef

`device_type`

device_type: DeviceKind

source

`from_device()`

static from_device(device)

source

Converts a Device or DeviceRef to a DeviceRef.

Parameters:: device (Device | DeviceRef)
Return type:: DeviceRef

`from_mlir()`

static from_mlir(attr)

source

Returns a device reference from an MLIR attribute.

Parameters:: attr (DeviceRefAttr)
Return type:: DeviceRef

`id`

id: int

source

`is_cpu()`

is_cpu()

source

Returns True if the device is a CPU device.

Return type:: bool

`is_gpu()`

is_gpu()

source

Returns True if the device is a GPU device.

Return type:: bool

`to_device()`

to_device()

source

Converts a device reference to a concrete driver Device.

Return type:: Device

`to_mlir()`

to_mlir()

source

Returns an MLIR attribute representing the device.

Return type:: DeviceRefAttr

DeviceRef​

CPU()​

GPU()​

device_type​

from_device()​

from_mlir()​

id​

is_cpu()​

is_gpu()​

to_device()​

to_mlir()​