Skip to content

declearn.utils.DevicePolicy

Dataclass to store parameters defining a device-selection policy.

This class merely defines a shared language of keyword-arguments to define whether to back computations on CPU or on a GPU device.

It is meant to be instantiated as a global variable that holds that information, and can be accessed by framework-specific backend code so as to take the required steps towards implementing that policy.

To access or update the current global DevicePolicy, please use the getter and setter functions: declearn.utils.get_device_policy and declearn.utils.set_device_policy.

Attributes:

Name Type Description
gpu bool

Whether to use a GPU device rather than the CPU one to back data and computations. If no GPU is available, use CPU with a warning.
idx int or None

Optional index of the GPU device to use. If None, select one arbitrarily. If this index exceeds the number of available GPUs, select one arbitrarily, with a warning.
Source code in declearn/utils/_device_policy.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
@dataclasses.dataclass
class DevicePolicy:
    """Dataclass to store parameters defining a device-selection policy.

    This class merely defines a shared language of keyword-arguments to
    define whether to back computations on CPU or on a GPU device.

    It is meant to be instantiated as a global variable that holds that
    information, and can be accessed by framework-specific backend code
    so as to take the required steps towards implementing that policy.

    To access or update the current global DevicePolicy, please use the
    getter and setter functions: `declearn.utils.get_device_policy` and
    `declearn.utils.set_device_policy`.

    Attributes
    ----------
    gpu: bool
        Whether to use a GPU device rather than the CPU one to back data
        and computations. If no GPU is available, use CPU with a warning.
    idx: int or None
        Optional index of the GPU device to use.
        If None, select one arbitrarily.
        If this index exceeds the number of available GPUs, select one
        arbitrarily, with a warning.
    """

    gpu: bool
    idx: Optional[int]

    def __post_init__(self) -> None:
        if not isinstance(self.gpu, bool):
            raise TypeError(
                f"DevicePolicy 'gpu' should be a bool, not '{type(self.gpu)}'."
            )
        if not (self.idx is None or isinstance(self.idx, int)):
            raise TypeError(
                "DevicePolicy 'idx' should be None or an int, not "
                f"'{type(self.idx)}'."
            )