The devices currently supported by Ivy are as follows:
In a similar manner to the
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.
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.
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#
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
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.
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:
deviceargument is provided, use this directly
otherwise, if an array is present in the arguments (very rare if the
deviceargument is present), set
arrto this array. This will then be used to infer the device by calling
ivy.dev()on the array
otherwise, if no arrays are present in the arguments (by far the most common case if the
deviceargument is present), then use the global default device, which currently can either be
tpu:idx. The default device is settable via
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
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.
def zeros( shape: Union[int, Sequence[int]], *, dtype: jnp.dtype, device: jaxlib.xla_extension.Device, ) -> JaxArray:
def zeros( shape: Union[int, Sequence[int]], *, dtype: np.dtype, device: str, ) -> np.ndarray:
def zeros( shape: Union[int, Sequence[int]], *, dtype: tf.DType, device: str, ) -> Tensor:
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.
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
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!