abracudabra.device.base

abracudabra.device.base#

Define the base device class.

Attributes#

DeviceType

The device type, e.g., "cpu" or "cuda".

DEVICE_TYPES

The supported device types.

Classes#

Device

A device with a name and index.

Module Contents#

abracudabra.device.base.DeviceType[source]#

The device type, e.g., "cpu" or "cuda".

abracudabra.device.base.DEVICE_TYPES: frozenset[DeviceType][source]#

The supported device types.

class abracudabra.device.base.Device[source]#

Bases: NamedTuple

A device with a name and index.

type: DeviceType[source]#

The device type, e.g., "cpu" or "cuda".

idx: int | None = None[source]#

The device index, e.g., 0 or None.

__str__()[source]#

Return the device name.

Return type:

str

classmethod validate(device, idx=None)[source]#

Return a device, validating the device type and index.

Parameters:

  • device (object) – The device type.

  • idx (object | None) – The optional device index.

Returns:

The device.

Return type:

Device

classmethod from_str(device, /)[source]#

Return a device from a string.

The string should be in the format "device[:idx]".

Examples

>>> Device.from_str("cpu")
Device(type="cpu", idx=None)
>>> Device.from_str("cuda:1")
Device(type="cuda", idx=1)
Parameters:

device (str)

Return type:

Device

classmethod parse(device, /)[source]#

Return a device from a string or device.

If the input is already a device, it is returned as is. Otherwise, the input is parsed as a string.

Parameters:

device (str | Device | torch.device) – The device or device string (e.g., "cpu" or "cuda:1").

Returns:

The device.

Return type:

Device

to_torch()[source]#

Return a torch device.

Examples

>>> Device("cpu", None).to_torch()
device(type='cpu')
>>> Device("cuda", 1).to_torch()
device(type='cuda', index=1)
Return type:

torch.device

Contents