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 their ivy.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 their ivy.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 their ivy.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 their ivy.NativeArray counterparts, calls the function with those arguments and then converts the ivy.NativeArray instances in the output back to ivy.Array. This wrapping function is heavily used because it enables achieving the objective of ensuring that every ivy function could accept an ivy.Array and return an ivy.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 the dtype 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 the out argument for a function, we prefer to use it as it’s a more efficient implementation of the out 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 the ivy namespace, just like calling a function with ivy.Array arguments instead. Thus, the function can be called by passing an ivy.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 their ivy.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 their ivy.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 their ivy.NativeShape counterparts, calls the function with those arguments and then converts the ivy.NativeShape instances in the output back to ivy.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