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.

Following are some of the wrapping functions currently used:

  1. 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.

  2. 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.

  3. 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.

  4. 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.

  5. 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.

  6. 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.

  7. 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.

  8. 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. Whenever there’s a ivy.Container argument, this wrapping function defers to the corresponding Containers static method to facilitate the same. As a result, the function can be called by passing an ivy.Container to any or all of its arguments.

  9. integer_array_to_float: This wrapping function enables conversion of integer array inputs in the positional and keyword arguments to a function to the default float dtype. This is currently used to support integer array arguments to functions for which one or more backend frameworks only non-integer numeric dtypes.

  10. handle_cmd_line_args: This wrapping function enables us to arbitrarily sample backend at test time using Hypothesis strategies. This enables us to infer the framework and generate appropriate data types directly inside the @given decorator. With this approach in place, it’s no longer necessary to check if the data type is supported and skip the test if it’s not. Another place wherein this decorator is helpful is when we perform configurable argument testing for the parameters (as_variable, with_out, native_array, container, instance_method, test_gradients) through the command line. The corresponding flags are used to set these values.

  11. handle_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 wrapper, it is necessary to set the mixed_function attribute of the function to True, as is done for example in ivy.linear. This attribute also automatically activates the inputs_to_ivy_arrays wrapper when calling the compositional implementation of the mixed function, and the outputs_to_ivy_arrays and inputs_to_native_arrays wrappers when calling the backend implementation. It is therefore preferable to avoid adding those decorators to mixed functions manually.

When calling _wrap_function during Backend Setting, firstly the attributes of the functions are checked to get all the wrapping functions for a particular functions. 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 channel or in the function wrapping forum!

Video