Function Wrapping#
When a backend framework is set by calling ivy.set_backend(backend_name)
, then all Ivy functions are wrapped.
This is achieved by calling _wrap_function, which will apply the appropriate wrapping to the given function, based on what decorators it has.
For example, abs has the decorators @to_native_arrays_and_back
and @handle_out_argument
, and so the backend implementations will also be wrapped with the to_native_arrays_and_back and handle_out_argument wrappers.
The new function returned by _wrap_function
is a replacement of the original function with extra code added to support requirements common to many functions in the API.
This is the main purpose of the wrapping, to avoid code duplication which would exist if we added identical logic in every single function independently.
Depending on the function being wrapped, the new function might handle Arrays, Inplace Updates, Data Types and/or Devices.
Our test decorators actually transforms to @given
decorators at Pytest collecting time, therefore this allows us to use other Hypothesis decorators like, @reproduce_failure
, @settings
, @seed
.
Decorator order#
The order in which Ivy decorators are applied is important. It is important to follow this order, as the functionality of many functions depends on it. If the decorators are applied in the wrong order, the test may fail or the function may not behave as expected. The following is the recommended order to follow :
@handle_complex_input
@infer_device
@handle_device_shifting
@infer_dtype
@handle_array_function
@outputs_to_ivy_arrays
@outputs_to_ivy_shapes
@outputs_to_native_arrays
@inputs_to_native_arrays
@inputs_to_native_shapes
@inputs_to_ivy_arrays
@handle_out_argument
@handle_view_indexing
@handle_view
@handle_array_like_without_promotion
@handle_partial_mixed_function
@handle_nestable
@handle_ragged
@handle_backend_invalid
@handle_exceptions
@handle_nans
This recommended order is followed to ensure that tests are efficient and accurate. It is important to follow this order because the decorators depend on each other. For example, the @infer_device
decorator needs to be applied before the @infer_dtype
decorator, because the @infer_dtype
decorator needs to know the device of the function in order to infer the data type.
Conversion Wrappers#
inputs_to_native_arrays : This wrapping function converts all
ivy.Array
instances in the arguments to theirivy.NativeArray
counterparts, based on the Backend Setting before calling the function.inputs_to_ivy_arrays : This wrapping function converts all
ivy.NativeArray
instances in the arguments to theirivy.Array
counterparts, based on the Backend Setting before calling the function.outputs_to_ivy_arrays : This wrapping function converts all
ivy.NativeArray
instances in the outputs to theirivy.Array
counterparts, based on the Backend Setting before calling the function.to_native_arrays_and_back : This wrapping function converts all
ivy.Array
instances in the arguments to theirivy.NativeArray
counterparts, calls the function with those arguments and then converts theivy.NativeArray
instances in the output back toivy.Array
. This wrapping function is heavily used because it enables achieving the objective of ensuring that every ivy function could accept anivy.Array
and return anivy.Array
, making it independent of the Backend Setting.
Inference Wrappers#
infer_dtype : This wrapping function infers the dtype argument to be passed to a function based on the array arguments passed to it. If
dtype
is explicitly passed to the function, then it is used directly. This wrapping function could be found in functions from the creation submodule such as zeros where we then allow the user to not enter thedtype
argument to such functions.infer_device : Similar to the infer_dtype wrapping function, the infer_device function wrapping infers the
device
argument to be passed to a function based on the first array argument passed to it. This wrapping function is also used a lot in functions from the creation submodule such as asarray, where we want to create the ivy.Array on the same device as the input array.
Out Argument Support#
handle_out_argument : This wrapping function is used in nearly all ivy functions. It enables appropriate handling of the
out
argument of functions. In cases where the backend framework natively supports theout
argument for a function, we prefer to use it as it’s a more efficient implementation of theout
argument for that particular backend framework. But in cases when it isn’t supported, we support it anyway with Inplace Updates.
Nestable Support#
handle_nestable : This wrapping function enables the use of
ivy.Container
arguments in functions and directly calling them through theivy
namespace, just like calling a function withivy.Array
arguments instead. Thus, the function can be called by passing anivy.Container
to any or all of its arguments.
Partial Mixed Function Support#
handle_partial_mixed_function: This wrapping function enables switching between compositional and primary implementations of Mixed Functions based on some condition on the arguments of the function.
The condition is specified through a lambda function which when evaluates to True the primary implementation is run and otherwise the compositional implementation is executed.
For backends that have a primary implementation of a mixed function, the reference to the compositional implementation is stored as an attribute inside the backend function during backend setting. To make use of this decorator, one must
add the
partial_mixed_handler
attribute containing the lambda function to the backend implementation. Here’s an example from the torch backend implementation of linear.
Shape Conversion#
inputs_to_native_shapes : This wrapping function converts all
ivy.Shape
instances in the arguments to theirivy.NativeShape
counterparts, based on the Backend Setting before calling the function.outputs_to_ivy_shapes : This wrapping function converts all
ivy.NativeShape
instances in the outputs to theirivy.Shape
counterparts, based on the Backend Setting before calling the function.to_native_shapes_and_back : This wrapping function converts all
ivy.Shape
instances in the arguments to theirivy.NativeShape
counterparts, calls the function with those arguments and then converts theivy.NativeShape
instances in the output back toivy.Shape
.
View Handling#
handle_view : This wrapping function performs view handling based on our Views policy.
handle_view_indexing : This wrapping function is aimed at handling views for indexing.
Exception Handling#
handle_exceptions : This wrapping function helps in catching native exceptions and unifying them into IvyException or the relevant subclasses. More information can be found in the Exception Handling section.
Miscellaneous Wrappers#
handle_array_function : This wrapping function enables Integrating custom classes with Ivy
handle_complex_input : This wrapping function enables handling of complex numbers. It introduces a keyword argument
complex_mode
, which is used to choose the function’s behaviour as per the wrapper’s docstring.
When calling _wrap_function during Backend Setting, firstly the attributes of the functions are checked to get all the wrapping functions for a particular function. Then all the wrapping functions applicable to a function are used to wrap the function.
Each of these topics and each associated piece of logic added by the various wrapper functions are covered in more detail in the next sections. For now, suffice it to say that they do quite a lot.
Round Up
This should have hopefully given you a good feel for how function wrapping is applied to functions in Ivy.
If you have any questions, please feel free to reach out on discord in the function wrapping thread!
Video