Ctrl+K

Function Arguments#

Here, we explain how the function arguments differ between the placeholder implementation at ivy/functional/ivy/category_name.py, and the backend-specific implementation at ivy/functional/backends/backend_name/category_name.py.

Many of these points are already addressed in the previous sections: Arrays, Data Types, Devices and Inplace Updates. However, we thought it would be convenient to revisit all of these considerations in a single section, dedicated to function arguments.

As for type-hints, all functions in the Ivy API at ivy/functional/ivy/category_name.py should have full and thorough type-hints. Likewise, all backend implementations at ivy/functional/backends/backend_name/category_name.py should also have full and thorough type-hints.

In order to understand the various requirements for function arguments, it’s useful to first look at some examples.

Examples#

For the purposes of explanation, we will use four functions as examples: ivy.tan(), ivy.roll(), ivy.add() and ivy.zeros().

We present both the Ivy API signature and also a backend-specific signature for each function:

# Ivy
@handle_exceptions
@handle_nestable
@handle_array_like_without_promotion
@handle_out_argument
@to_native_arrays_and_back
@handle_array_function
def tan(
    x: Union[ivy.Array, ivy.NativeArray],
    /,
    *,
    out: Optional[ivy.Array] = None
) -> ivy.Array:

# PyTorch
@handle_numpy_arrays_in_specific_backend
def tan(
    x: torch.Tensor,
    /,
    *,
    out: Optional[torch.Tensor] = None
) -> torch.Tensor:

# Ivy
@handle_exceptions
@handle_nestable
@handle_array_like_without_promotion
@handle_out_argument
@to_native_arrays_and_back
@handle_array_function
def roll(
    x: Union[ivy.Array, ivy.NativeArray],
    /,
    shift: Union[int, Sequence[int]],
    *,
    axis: Optional[Union[int, Sequence[int]]] = None,
    out: Optional[ivy.Array] = None,
) -> ivy.Array:

# NumPy
def roll(
    x: np.ndarray,
    /,
    shift: Union[int, Sequence[int]],
    *,
    axis: Optional[Union[int, Sequence[int]]] = None,
    out: Optional[np.ndarray] = None,
) -> np.ndarray:

# Ivy
@handle_exceptions
@handle_nestable
@handle_out_argument
@to_native_arrays_and_back
@handle_array_function
def add(
    x1: Union[float, ivy.Array, ivy.NativeArray],
    x2: Union[float, ivy.Array, ivy.NativeArray],
    /,
    *,
    alpha: Optional[Union[int, float]] = None,
    out: Optional[ivy.Array] = None,
) -> ivy.Array:

# TensorFlow
def add(
    x1: Union[float, tf.Tensor, tf.Variable],
    x2: Union[float, tf.Tensor, tf.Variable],
    /,
    *,
    alpha: Optional[Union[int, float]] = None,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:

# Ivy
@handle_nestable
@handle_array_like_without_promotion
@handle_out_argument
@inputs_to_native_shapes
@outputs_to_ivy_arrays
@handle_array_function
@infer_dtype
@infer_device
def zeros(
    shape: Union[ivy.Shape, ivy.NativeShape],
    *,
    dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None,
    device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None,
    out: Optional[ivy.Array] = None
) -> ivy.Array:

# JAX
def zeros(
    shape:  Union[ivy.NativeShape, Sequence[int]],
    *,
    dtype: jnp.dtype,
    device: jaxlib.xla_extension.Device,
    out: Optional[JaxArray] = None,
) -> JaxArray:

Positional and Keyword Arguments#

In both signatures, we follow the Array API Standard convention about positional and keyword arguments.

  • Positional parameters must be positional-only parameters. Positional-only parameters have no externally-usable name. When a method accepting positional-only parameters is called, positional arguments are mapped to these parameters based solely on their order. This is indicated with an / after all the position-only arguments.

  • Optional parameters must be keyword-only arguments. A * must be added before any of the keyword-only arguments.

Nearly all the functions in the Array API Standard convention have strictly positional-only and keyword-only arguments, with an exception of few creation functions such as ones(shape, *, dtype=None, device=None) , linspace(start, stop, /, num, *, dtype=None, device=None, endpoint=True) etc. The rationale behind this is purely a convention. The shape argument is often passed as a keyword, while the num argument in linspace is often passed as a keyword for improved understandability of the code. Therefore, given that Ivy fully adheres to the Array API Standard, Ivy also adopts these same exceptions to the general rule for the shape and num arguments in these functions.

Input Arrays#

In each example, we can see that the input arrays have type Union[ivy.Array, ivy.NativeArray] whereas the output arrays have type ivy.Array. This is the case for all functions in the Ivy API. We always return an ivy.Array instance to ensure that any subsequent Ivy code is fully framework-agnostic, with all operators performed on the returned array now handled by the special methods of the ivy.Array class, and not the special methods of the backend array class (ivy.NativeArray). For example, calling any of (+, -, *, / etc.) on the array will result in (__add__, __sub__, __mul__, __div__ etc.) being called on the array class.

ivy.NativeArray instances are also not permitted for the out argument, which is used in many functions. This is because the out argument dictates the array to which the result should be written, and so it effectively serves the same purpose as the function return when no out argument is specified. This is all explained in more detail in the Arrays section.

out Argument#

The out argument should always be provided as a keyword-only argument, and it should be added to all functions in the Ivy API and backend API which support inplace updates, with a default value of None in all cases. The out argument is explained in more detail in the Inplace Updates section.

dtype and device arguments#

In the Ivy API at ivy/functional/ivy/category_name.py, the dtype and device arguments should both always be provided as keyword-only arguments, with a default value of None. In contrast, these arguments should both be added as required arguments in the backend implementation at ivy/functional/backends/backend_name/category_name.py. In a nutshell, by the time the backend implementation is entered, the correct dtype and device to use have both already been correctly handled by code which is wrapped around the backend implementation. This is further explained in the Data Types and Devices sections respectively.

Numbers in Operator Functions#

All operator functions (which have a corresponding such as +, -, *, /) must also be fully compatible with numbers (float or int) passed into any of the array inputs, even in the absence of any arrays. For example, ivy.add(1, 2), ivy.add(1.5, 2) and ivy.add(1.5, ivy.array([2])) should all run without error. Therefore, the type hints for ivy.add() include float as one of the types in the Union for the array inputs, and also as one of the types in the Union for the output. PEP 484 Type Hints states that “when an argument is annotated as having type float, an argument of type int is acceptable”. Therefore, we only include float in the type hints.

Integer Sequences#

For sequences of integers, generally the Array API Standard dictates that these should be of type Tuple[int], and not List[int]. However, in order to make Ivy code less brittle, we accept arbitrary integer sequences Sequence[int] for such arguments (which includes list, tuple etc.). This does not break the standard, as the standard is only intended to define a subset of required behaviour. The standard can be freely extended, as we are doing here. Good examples of this are the axis argument of ivy.roll() and the shape argument of ivy.zeros(), as shown above.

Nestable Functions#

Most functions in the Ivy API can also consume and return ivy.Container instances in place of the any of the function arguments. If an ivy.Container is passed, then the function is mapped across all of the leaves of this container. Because of this feature, we refer to these functions as nestable functions. However, because so many functions in the Ivy API are indeed nestable functions, and because this flexibility applies to every argument in the function, every type hint for these functions should technically be extended like so: Union[original_type, ivy.Container].

However, this would be very cumbersome, and would only serve to hinder the readability of the docs. Therefore, we simply omit these ivy.Container type hints from nestable functions, and instead mention in the docstring whether the function is nestable or not.

Round Up

These examples should hopefully give you a good understanding of what is required when adding function arguments.

If you have any questions, please feel free to reach out on discord in the function arguments thread!

Video

On this page