Devices#
The devices currently supported by Ivy are as follows:
cpu
gpu:idx
tpu:idx
In a similar manner to the ivy.Dtype
and ivy.NativeDtype
classes (see Data Types), there is both an ivy.Device class and an ivy.NativeDevice
class, with ivy.NativeDevice
initially set as an empty class.
The ivy.Device
class derives from str
, and has simple logic in the constructor to verify that the string formatting is correct.
When a backend is set, the ivy.NativeDtype
is replaced with the backend-specific device class.
Device Module#
The device.py module provides a variety of functions for working with devices.
A few examples include ivy.get_all_ivy_arrays_on_dev()
which gets all arrays which are currently alive on the specified device, ivy.dev()
which gets the device for input array, and ivy.num_gpus()
which determines the number of available GPUs for use with the backend framework.
Many functions in the device.py
module are convenience functions, which means that they do not directly modify arrays, as explained in the Function Types section.
For example, the following are all convenience functions: ivy.total_mem_on_dev, which gets the total amount of memory for a given device, ivy.dev_util, which gets the current utilization (%) for a given device, ivy.num_cpu_cores, which determines the number of cores available in the CPU, and ivy.default_device, which returns the correct device to use.
ivy.default_device is arguably the most important function.
Any function in the functional API that receives a device
argument will make use of this function, as explained below.
Arguments in other Functions#
Like with dtype
, all device
arguments are also keyword-only.
All creation functions include the device
argument, for specifying the device on which to place the created array.
Some other functions outside of the creation.py
submodule also support the device
argument, such as ivy.random_uniform()
which is located in random.py
, but this is simply because of dual categorization.
ivy.random_uniform()
is also essentially a creation function, despite not being located in creation.py
.
The device
argument is generally not included for functions which accept arrays in the input and perform operations on these arrays.
In such cases, the device of the output arrays is the same as the device for the input arrays.
In cases where the input arrays are located on different devices, an error will generally be thrown, unless the function is specific to distributed training.
The device
argument is handled in infer_device for all functions which have the @infer_device
decorator, similar to how dtype
is handled.
This function calls ivy.default_device in order to determine the correct device.
As discussed in the Function Wrapping section, this is applied to all applicable functions dynamically during backend setting.
Overall, ivy.default_device infers the device as follows:
if the
device
argument is provided, use this directlyotherwise, if an array is present in the arguments (very rare if the
device
argument is present), setarr
to this array. This will then be used to infer the device by callingivy.dev()
on the arrayotherwise, if no arrays are present in the arguments (by far the most common case if the
device
argument is present), then use the global default device, which currently can either becpu
,gpu:idx
ortpu:idx
. The default device is settable viaivy.set_default_device()
.
For the majority of functions which defer to infer_device for handling the device, these steps will have been followed and the device
argument will be populated with the correct value before the backend-specific implementation is even entered into.
Therefore, whereas the device
argument is listed as optional in the ivy API at ivy/functional/ivy/category_name.py
, the argument is listed as required in the backend-specific implementations at ivy/functional/backends/backend_name/category_name.py
.
This is exactly the same as with the dtype
argument, as explained in the Data Types section.
Let’s take a look at the function ivy.zeros()
as an example.
The implementation in ivy/functional/ivy/creation.py
has the following signature:
@outputs_to_ivy_arrays
@handle_out_argument
@infer_dtype
@infer_device
def zeros(
shape: Union[int, Sequence[int]],
*,
dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None,
device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None,
) -> ivy.Array:
Whereas the backend-specific implementations in ivy/functional/backends/backend_name/creation.py
all list device
as required.
Jax:
def zeros(
shape: Union[int, Sequence[int]],
*,
dtype: jnp.dtype,
device: jaxlib.xla_extension.Device,
) -> JaxArray:
NumPy:
def zeros(
shape: Union[int, Sequence[int]],
*,
dtype: np.dtype,
device: str,
) -> np.ndarray:
TensorFlow:
def zeros(
shape: Union[int, Sequence[int]],
*,
dtype: tf.DType,
device: str,
) -> Tensor:
PyTorch:
def zeros(
shape: Union[int, Sequence[int]],
*,
dtype: torch.dtype,
device: torch.device,
) -> Tensor:
This makes it clear that these backend-specific functions are only enterred into once the correct device
has been determined.
However, the device
argument for functions without the @infer_device
decorator is not handled by infer_device, and so these defaults must be handled by the backend-specific implementations themselves, by calling ivy.default_device()
internally.
Round Up
This should have hopefully given you a good feel for devices, and how these are handled in Ivy.
If you have any questions, please feel free to reach out on discord in the devices channel or in the devices forum!
Video