From 8847c9b2af3837fb25f416439377646e8243f7d2 Mon Sep 17 00:00:00 2001 From: Adlai-1 Date: Tue, 6 Feb 2024 17:36:19 +0000 Subject: [PATCH 1/2] made dcostring changes and added new examples to it. --- ivy/functional/ivy/device.py | 177 ++++++++++++++++++++++------------- 1 file changed, 111 insertions(+), 66 deletions(-) diff --git a/ivy/functional/ivy/device.py b/ivy/functional/ivy/device.py index 9187040ad1ffc..2848f19e1f932 100644 --- a/ivy/functional/ivy/device.py +++ b/ivy/functional/ivy/device.py @@ -59,7 +59,8 @@ def __init__( device: Union[ivy.Device, ivy.NativeDevice], /, ) -> None: - """Initialize the DefaultDevice class. + """ + Initialize the DefaultDevice class. Parameters ---------- @@ -75,7 +76,8 @@ def __init__( self._dev = device def __enter__(self): - """Enter the runtime context related to the specified device. + """ + Enter the runtime context related to the specified device. Returns ------- @@ -101,7 +103,8 @@ def __exit__( exc_val: Optional[Type[BaseException]], exc_tb: Optional[types.TracebackType], ) -> Union[ivy.Device, str]: - """Exit the runtime context related to the specified device. + """ + Exit the runtime context related to the specified device. Parameters ---------- @@ -175,7 +178,8 @@ def get_all_ivy_arrays_on_dev( device: Union[ivy.Device, ivy.NativeDevice], /, ) -> ivy.Container: - """Get all ivy arrays which are currently alive on the specified device. + """ + Get all ivy arrays which are currently alive on the specified device. Parameters ---------- @@ -210,8 +214,8 @@ def get_all_ivy_arrays_on_dev( @handle_exceptions def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: - """Return the number of arrays which are currently alive on the specified - device. + """ + Return the number of arrays which are currently alive on the specified device. Parameters ---------- @@ -252,8 +256,9 @@ def print_all_ivy_arrays_on_dev( device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, attr_only: bool = True, ) -> None: - """Print the shape and dtype for all ivy arrays which are currently alive - on the specified device. + """ + Print the shape and dtype for all ivy arrays which are currently alive on the + specified device. Parameters ---------- @@ -290,8 +295,9 @@ def print_all_ivy_arrays_on_dev( @handle_exceptions def set_soft_device_mode(mode: bool) -> None: - """Set the mode of whether to move input arrays to `ivy.default_device()` - before performing an operation. + """ + Set the mode of whether to move input arrays to `ivy.default_device()` before + performing an operation. Parameter --------- @@ -315,8 +321,9 @@ def set_soft_device_mode(mode: bool) -> None: @handle_exceptions def unset_soft_device_mode() -> None: - """Reset the mode of moving input arrays to `ivy.default_device()` before - performing an operation. + """ + Reset the mode of moving input arrays to `ivy.default_device()` before performing an + operation. Examples -------- @@ -344,7 +351,8 @@ def unset_soft_device_mode() -> None: def dev( x: Union[ivy.Array, ivy.NativeArray], /, *, as_native: bool = False ) -> Union[ivy.Device, ivy.NativeDevice]: - """Get the native device handle for input array x. + """ + Get the native device handle for input array x. Parameters ---------- @@ -382,7 +390,8 @@ def dev( @handle_exceptions def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: - """Convert device to string representation. + """ + Convert device to string representation. Parameters ---------- @@ -405,7 +414,8 @@ def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: @handle_exceptions def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: - """Convert device string representation to native device type. + """ + Convert device string representation to native device type. Parameters ---------- @@ -449,7 +459,8 @@ def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeD @handle_exceptions def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: - """Clear memory cache on target device. + """ + Clear memory cache on target device. Parameters ---------- @@ -468,8 +479,9 @@ def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> N @handle_exceptions def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: - """Get the total amount of memory (in GB) for a given device string. In - case of CPU, the total RAM is returned. + """ + Get the total amount of memory (in GB) for a given device string. In case of CPU, + the total RAM is returned. Parameters ---------- @@ -511,8 +523,9 @@ def used_mem_on_dev( *, process_specific: bool = False, ) -> float: - """Get the used memory (in GB) for a given device string. In case of CPU, - the used RAM is returned. + """ + Get the used memory (in GB) for a given device string. In case of CPU, the used RAM + is returned. Parameters ---------- @@ -570,8 +583,9 @@ def percent_used_mem_on_dev( *, process_specific: bool = False, ) -> float: - """Get the percentage used memory for a given device string. In case of - CPU, the used RAM is returned. + """ + Get the percentage used memory for a given device string. In case of CPU, the used + RAM is returned. Parameters ---------- @@ -631,7 +645,8 @@ def dev_util( device: Union[ivy.Device, ivy.NativeDevice], /, ) -> float: - """Get the current utilization (%) for a given device. + """ + Get the current utilization (%) for a given device. Parameters ---------- @@ -673,7 +688,8 @@ def dev_util( @handle_exceptions def gpu_is_available() -> bool: - """Determine whether a GPU is available to use, with the backend framework. + """ + Determine whether a GPU is available to use, with the backend framework. Returns ------- @@ -690,7 +706,8 @@ def gpu_is_available() -> bool: @handle_exceptions def num_cpu_cores(*, logical: bool = True) -> int: - """Determine the number of cores available in the cpu. + """ + Determine the number of cores available in the cpu. Parameters ---------- @@ -715,7 +732,8 @@ def num_cpu_cores(*, logical: bool = True) -> int: @handle_exceptions def num_gpus() -> int: - """Determine the number of available GPUs, with the backend framework. + """ + Determine the number of available GPUs, with the backend framework. Returns ------- @@ -732,7 +750,8 @@ def num_gpus() -> int: @handle_exceptions def tpu_is_available() -> bool: - """Determine whether a TPU is available to use, with the backend framework. + """ + Determine whether a TPU is available to use, with the backend framework. Returns ------- @@ -760,11 +779,12 @@ def default_device( item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, as_native: Optional[bool] = None, ) -> Union[ivy.Device, ivy.NativeDevice]: - """Return the input device or the default device. If the as_native flag is - set, the device will be converted to a native device. If the item is - provided, the item's device is returned. If the device is not provided, the - last default device is returned. If a default device has not been set, the - first gpu is returned if available, otherwise the cpu is returned. + """ + Return the input device or the default device. If the as_native flag is set, the + device will be converted to a native device. If the item is provided, the item's + device is returned. If the device is not provided, the last default device is + returned. If a default device has not been set, the first gpu is returned if + available, otherwise the cpu is returned. Parameters ---------- @@ -826,30 +846,48 @@ def default_device( @handle_exceptions def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: - """Set the default device to given device instance. + """ + Sets the default device to the argument provided in the function. Parameters ---------- device - The device to set as the default device + The device to be set as the default device. + + Returns + ------- + ret + The new default device. Examples -------- - >>> ivy.set_default_device("cpu") >>> ivy.default_device() 'cpu' - >>> ivy.set_backend("torch") - >>> ivy.set_default_device("gpu:0") - >>> ivy.default_device(as_native=True) - device(type='cuda', index=0) + >>> ivy.set_backend('jax') + >>> ivy.set_default_device('gpu:0') + >>> ivy.default_device() + 'gpu:0' - >>> import torch - >>> ivy.set_backend("torch") - >>> device = torch.device("cuda") - >>> ivy.set_default_device(device) - >>> ivy.default_device(as_native=True) - device(type='cuda') + >>> ivy.set_backend('torch') + >>> ivy.set_default_device('gpu:1') + >>> ivy.default_device() + 'gpu:1 + + >>> ivy.set_backend('tensorflow') + >>> ivy.set_default_device('tpu:0) + >>> ivy.default_device() + 'tpu:0 + + >>> ivy.set_backend('paddle') + >>> ivy.set_default_device('cpu) + >>> ivy.default_device() + 'cpu' + + >>> ivy.set_backend('mxnet') + >>> ivy.set_default_device('cpu') + >>> ivy.default_device() + 'cpu' """ global default_device_stack default_device_stack.append(device) @@ -857,7 +895,8 @@ def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: @handle_exceptions def unset_default_device() -> None: - """Reset the default device to "cpu". + """ + Reset the default device to "cpu". Examples -------- @@ -890,8 +929,8 @@ def to_device( stream: Optional[Union[int, Any]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """Move the input array x to the desired device, specified by device - string. + """ + Move the input array x to the desired device, specified by device string. Parameters ---------- @@ -931,8 +970,9 @@ def split_factor( device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, /, ) -> float: - """Get a device's global split factor, which can be used to scale the - device's batch splitting chunk sizes across the codebase. + """ + Get a device's global split factor, which can be used to scale the device's batch + splitting chunk sizes across the codebase. If the global split factor is set for a given device, returns the split factor value for the device from the split factors dictionary @@ -968,8 +1008,9 @@ def split_factor( def set_split_factor( factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None ) -> None: - """Set the global split factor for a given device, which can be used to - scale batch splitting chunk sizes for the device across the codebase. + """ + Set the global split factor for a given device, which can be used to scale batch + splitting chunk sizes for the device across the codebase. Parameters ---------- @@ -1026,9 +1067,10 @@ def split_func_call( stop_gradients: bool = False, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, ) -> Union[ivy.Array, ivy.NativeArray]: - """Call a function by splitting its inputs along a given axis, and calling - the function in chunks, rather than feeding the entire input array at once. - This can be useful to reduce memory usage of the device the arrays are on. + """ + Call a function by splitting its inputs along a given axis, and calling the function + in chunks, rather than feeding the entire input array at once. This can be useful to + reduce memory usage of the device the arrays are on. Parameters ---------- @@ -1201,10 +1243,10 @@ def _get_devices(fn: Callable, complement: bool = True) -> Tuple: def function_supported_devices( fn: Callable, recurse: bool = True ) -> Union[Tuple, dict]: - """Return the supported devices of the current backend's function. The - function returns a dict containing the supported devices for the - compositional and primary implementations in case of partial mixed - functions. + """ + Return the supported devices of the current backend's function. The function returns + a dict containing the supported devices for the compositional and primary + implementations in case of partial mixed functions. Parameters ---------- @@ -1253,10 +1295,10 @@ def function_supported_devices( def function_unsupported_devices( fn: Callable, recurse: bool = True ) -> Union[Tuple, dict]: - """Return the unsupported devices of the current backend's function. The - function returns a dict containing the unsupported devices for the - compositional and primary implementations in case of partial mixed - functions. + """ + Return the unsupported devices of the current backend's function. The function + returns a dict containing the unsupported devices for the compositional and primary + implementations in case of partial mixed functions. Parameters ---------- @@ -1302,7 +1344,8 @@ def function_unsupported_devices( class Profiler(abc.ABC): - """The profiler class is used to profile the execution of some code. + """ + The profiler class is used to profile the execution of some code. Parameters ---------- @@ -1315,7 +1358,8 @@ def __init__(self, save_dir: str): @abc.abstractmethod def start(self): - """Start the profiler. + """ + Start the profiler. This should be called before the code to be profiled. """ @@ -1323,7 +1367,8 @@ def start(self): @abc.abstractmethod def stop(self): - """Stop the profiler. + """ + Stop the profiler. This should be called after the code to be profiled. """ From 11a692d9a9a301dfec0dd589113dd35ad63b4ed2 Mon Sep 17 00:00:00 2001 From: ivy-branch Date: Tue, 6 Feb 2024 17:47:33 +0000 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=A4=96=20Lint=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ivy/functional/ivy/device.py | 136 ++++++++++++++--------------------- 1 file changed, 54 insertions(+), 82 deletions(-) diff --git a/ivy/functional/ivy/device.py b/ivy/functional/ivy/device.py index 2848f19e1f932..b8a50593edb6e 100644 --- a/ivy/functional/ivy/device.py +++ b/ivy/functional/ivy/device.py @@ -59,8 +59,7 @@ def __init__( device: Union[ivy.Device, ivy.NativeDevice], /, ) -> None: - """ - Initialize the DefaultDevice class. + """Initialize the DefaultDevice class. Parameters ---------- @@ -76,8 +75,7 @@ def __init__( self._dev = device def __enter__(self): - """ - Enter the runtime context related to the specified device. + """Enter the runtime context related to the specified device. Returns ------- @@ -103,8 +101,7 @@ def __exit__( exc_val: Optional[Type[BaseException]], exc_tb: Optional[types.TracebackType], ) -> Union[ivy.Device, str]: - """ - Exit the runtime context related to the specified device. + """Exit the runtime context related to the specified device. Parameters ---------- @@ -178,8 +175,7 @@ def get_all_ivy_arrays_on_dev( device: Union[ivy.Device, ivy.NativeDevice], /, ) -> ivy.Container: - """ - Get all ivy arrays which are currently alive on the specified device. + """Get all ivy arrays which are currently alive on the specified device. Parameters ---------- @@ -214,8 +210,8 @@ def get_all_ivy_arrays_on_dev( @handle_exceptions def num_ivy_arrays_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> int: - """ - Return the number of arrays which are currently alive on the specified device. + """Return the number of arrays which are currently alive on the specified + device. Parameters ---------- @@ -256,9 +252,8 @@ def print_all_ivy_arrays_on_dev( device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, attr_only: bool = True, ) -> None: - """ - Print the shape and dtype for all ivy arrays which are currently alive on the - specified device. + """Print the shape and dtype for all ivy arrays which are currently alive + on the specified device. Parameters ---------- @@ -295,9 +290,8 @@ def print_all_ivy_arrays_on_dev( @handle_exceptions def set_soft_device_mode(mode: bool) -> None: - """ - Set the mode of whether to move input arrays to `ivy.default_device()` before - performing an operation. + """Set the mode of whether to move input arrays to `ivy.default_device()` + before performing an operation. Parameter --------- @@ -321,9 +315,8 @@ def set_soft_device_mode(mode: bool) -> None: @handle_exceptions def unset_soft_device_mode() -> None: - """ - Reset the mode of moving input arrays to `ivy.default_device()` before performing an - operation. + """Reset the mode of moving input arrays to `ivy.default_device()` before + performing an operation. Examples -------- @@ -351,8 +344,7 @@ def unset_soft_device_mode() -> None: def dev( x: Union[ivy.Array, ivy.NativeArray], /, *, as_native: bool = False ) -> Union[ivy.Device, ivy.NativeDevice]: - """ - Get the native device handle for input array x. + """Get the native device handle for input array x. Parameters ---------- @@ -390,8 +382,7 @@ def dev( @handle_exceptions def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: - """ - Convert device to string representation. + """Convert device to string representation. Parameters ---------- @@ -414,8 +405,7 @@ def as_ivy_dev(device: Union[ivy.Device, str], /) -> ivy.Device: @handle_exceptions def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeDevice: - """ - Convert device string representation to native device type. + """Convert device string representation to native device type. Parameters ---------- @@ -459,8 +449,7 @@ def as_native_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> ivy.NativeD @handle_exceptions def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: - """ - Clear memory cache on target device. + """Clear memory cache on target device. Parameters ---------- @@ -479,9 +468,8 @@ def clear_cached_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> N @handle_exceptions def total_mem_on_dev(device: Union[ivy.Device, ivy.NativeDevice], /) -> float: - """ - Get the total amount of memory (in GB) for a given device string. In case of CPU, - the total RAM is returned. + """Get the total amount of memory (in GB) for a given device string. In + case of CPU, the total RAM is returned. Parameters ---------- @@ -523,9 +511,8 @@ def used_mem_on_dev( *, process_specific: bool = False, ) -> float: - """ - Get the used memory (in GB) for a given device string. In case of CPU, the used RAM - is returned. + """Get the used memory (in GB) for a given device string. In case of CPU, + the used RAM is returned. Parameters ---------- @@ -583,9 +570,8 @@ def percent_used_mem_on_dev( *, process_specific: bool = False, ) -> float: - """ - Get the percentage used memory for a given device string. In case of CPU, the used - RAM is returned. + """Get the percentage used memory for a given device string. In case of + CPU, the used RAM is returned. Parameters ---------- @@ -645,8 +631,7 @@ def dev_util( device: Union[ivy.Device, ivy.NativeDevice], /, ) -> float: - """ - Get the current utilization (%) for a given device. + """Get the current utilization (%) for a given device. Parameters ---------- @@ -688,8 +673,7 @@ def dev_util( @handle_exceptions def gpu_is_available() -> bool: - """ - Determine whether a GPU is available to use, with the backend framework. + """Determine whether a GPU is available to use, with the backend framework. Returns ------- @@ -706,8 +690,7 @@ def gpu_is_available() -> bool: @handle_exceptions def num_cpu_cores(*, logical: bool = True) -> int: - """ - Determine the number of cores available in the cpu. + """Determine the number of cores available in the cpu. Parameters ---------- @@ -732,8 +715,7 @@ def num_cpu_cores(*, logical: bool = True) -> int: @handle_exceptions def num_gpus() -> int: - """ - Determine the number of available GPUs, with the backend framework. + """Determine the number of available GPUs, with the backend framework. Returns ------- @@ -750,8 +732,7 @@ def num_gpus() -> int: @handle_exceptions def tpu_is_available() -> bool: - """ - Determine whether a TPU is available to use, with the backend framework. + """Determine whether a TPU is available to use, with the backend framework. Returns ------- @@ -779,12 +760,11 @@ def default_device( item: Optional[Union[list, tuple, dict, ivy.Array, ivy.NativeArray]] = None, as_native: Optional[bool] = None, ) -> Union[ivy.Device, ivy.NativeDevice]: - """ - Return the input device or the default device. If the as_native flag is set, the - device will be converted to a native device. If the item is provided, the item's - device is returned. If the device is not provided, the last default device is - returned. If a default device has not been set, the first gpu is returned if - available, otherwise the cpu is returned. + """Return the input device or the default device. If the as_native flag is + set, the device will be converted to a native device. If the item is + provided, the item's device is returned. If the device is not provided, the + last default device is returned. If a default device has not been set, the + first gpu is returned if available, otherwise the cpu is returned. Parameters ---------- @@ -846,8 +826,7 @@ def default_device( @handle_exceptions def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: - """ - Sets the default device to the argument provided in the function. + """Sets the default device to the argument provided in the function. Parameters ---------- @@ -895,8 +874,7 @@ def set_default_device(device: Union[ivy.Device, ivy.NativeDevice], /) -> None: @handle_exceptions def unset_default_device() -> None: - """ - Reset the default device to "cpu". + """Reset the default device to "cpu". Examples -------- @@ -929,8 +907,8 @@ def to_device( stream: Optional[Union[int, Any]] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Move the input array x to the desired device, specified by device string. + """Move the input array x to the desired device, specified by device + string. Parameters ---------- @@ -970,9 +948,8 @@ def split_factor( device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, /, ) -> float: - """ - Get a device's global split factor, which can be used to scale the device's batch - splitting chunk sizes across the codebase. + """Get a device's global split factor, which can be used to scale the + device's batch splitting chunk sizes across the codebase. If the global split factor is set for a given device, returns the split factor value for the device from the split factors dictionary @@ -1008,9 +985,8 @@ def split_factor( def set_split_factor( factor: float, /, *, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None ) -> None: - """ - Set the global split factor for a given device, which can be used to scale batch - splitting chunk sizes for the device across the codebase. + """Set the global split factor for a given device, which can be used to + scale batch splitting chunk sizes for the device across the codebase. Parameters ---------- @@ -1067,10 +1043,9 @@ def split_func_call( stop_gradients: bool = False, device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None, ) -> Union[ivy.Array, ivy.NativeArray]: - """ - Call a function by splitting its inputs along a given axis, and calling the function - in chunks, rather than feeding the entire input array at once. This can be useful to - reduce memory usage of the device the arrays are on. + """Call a function by splitting its inputs along a given axis, and calling + the function in chunks, rather than feeding the entire input array at once. + This can be useful to reduce memory usage of the device the arrays are on. Parameters ---------- @@ -1243,10 +1218,10 @@ def _get_devices(fn: Callable, complement: bool = True) -> Tuple: def function_supported_devices( fn: Callable, recurse: bool = True ) -> Union[Tuple, dict]: - """ - Return the supported devices of the current backend's function. The function returns - a dict containing the supported devices for the compositional and primary - implementations in case of partial mixed functions. + """Return the supported devices of the current backend's function. The + function returns a dict containing the supported devices for the + compositional and primary implementations in case of partial mixed + functions. Parameters ---------- @@ -1295,10 +1270,10 @@ def function_supported_devices( def function_unsupported_devices( fn: Callable, recurse: bool = True ) -> Union[Tuple, dict]: - """ - Return the unsupported devices of the current backend's function. The function - returns a dict containing the unsupported devices for the compositional and primary - implementations in case of partial mixed functions. + """Return the unsupported devices of the current backend's function. The + function returns a dict containing the unsupported devices for the + compositional and primary implementations in case of partial mixed + functions. Parameters ---------- @@ -1344,8 +1319,7 @@ def function_unsupported_devices( class Profiler(abc.ABC): - """ - The profiler class is used to profile the execution of some code. + """The profiler class is used to profile the execution of some code. Parameters ---------- @@ -1358,8 +1332,7 @@ def __init__(self, save_dir: str): @abc.abstractmethod def start(self): - """ - Start the profiler. + """Start the profiler. This should be called before the code to be profiled. """ @@ -1367,8 +1340,7 @@ def start(self): @abc.abstractmethod def stop(self): - """ - Stop the profiler. + """Stop the profiler. This should be called after the code to be profiled. """