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:

  1. if the device argument is provided, use this directly

  2. otherwise, if an array is present in the arguments (very rare if the device argument is present), set arr to this array. This will then be used to infer the device by calling ivy.dev() on the array

  3. otherwise, 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 be cpu, gpu:idx or tpu:idx. The default device is settable via ivy.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