Max Interop#

Shorthands for utitly functions from MAX

nabla.utils.max_interop.accelerator(device_id=0)[source]#

Create an Accelerator device instance with the specified GPU ID.

Parameters:

device_id (int) – GPU ID (default is 0)

Returns:

An instance of the Accelerator class for the specified GPU.

Return type:

max.driver.Device

nabla.utils.max_interop.cpu()[source]#

Create a CPU device instance.

Returns:

An instance of the CPU class.

Return type:

max.driver.Device

nabla.utils.max_interop.device(device_name)[source]#

Get a device instance based on the provided device name.

Parameters:

device_name (str) – Name of the device (e.g., “cpu”, “cuda”, “mps”)

Returns:

An instance of the corresponding Device class.

Return type:

max.driver.Device

nabla.utils.max_interop.device_ref(device)[source]#

Get a DeviceRef instance for the specified device.

Parameters:

device (max.driver.Device) – The Device instance to reference.

Returns:

A DeviceRef instance for the specified device.

Return type:

<MagicMock id=’140511746857424’>

nabla.utils.max_interop.accelerator_count()[source]#

Get the number of available accelerators (GPUs).

Returns:

The number of available accelerators.

Return type:

int