pytme 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

tme/backends/matching_backend.py

@@ -6,11 +6,22 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Tuple, Callable, List
 from multiprocessing import shared_memory
+from typing import Tuple, Callable, List, Any, Union, Optional, Generator
 
-from numpy.typing import NDArray
-from ..types import ArrayLike, Scalar
+from ..types import BackendArray, NDArray, Scalar, shm_type
+
+
+def _create_metafunction(func_name: str) -> Callable:
+    """
+    Returns a wrapper of ``self._array_backend.func_name``.
+    """
+
+    def metafunction(self, *args, **kwargs) -> Any:
+        backend_func = getattr(self._array_backend, func_name)
+        return backend_func(*args, **kwargs)
+
+    return metafunction
 
 
 class MatchingBackend(ABC):
@@ -29,54 +40,69 @@ class MatchingBackend(ABC):
     ----------
     array_backend : object
         The backend object providing array functionalities.
-    default_dtype : type
-        Data type of real array instances, e.g. np.float32.
+    float_dtype : type
+        Data type of float array instances, e.g. np.float32.
     complex_dtype : type
         Data type of complex array instances, e.g. np.complex64.
-    default_dtype_int : type
+    int_dtype : type
         Data type of integer array instances, e.g. np.int32.
+    overflow_safe_dtype : type
+        Data type than can be used in reduction operations to avoid overflows.
 
     Attributes
     ----------
     _array_backend : object
-        The backend object used to delegate method and attribute calls.
-    _default_dtype : type
-        Data type of real array instances, e.g. np.float32.
+        The backend object providing array functionalities.
+    _float_dtype : type
+        Data type of float array instances, e.g. np.float32.
     _complex_dtype : type
         Data type of complex array instances, e.g. np.complex64.
-    _default_dtype_int : type
+    _int_dtype : type
         Data type of integer array instances, e.g. np.int32.
+    _overflow_safe_dtype : type
+        Data type than can be used in reduction operations to avoid overflows.
+    _fundamental_dtypes : Dict
+        Maps int, float and cmoplex python types to backend specific data types.
 
     Examples
     --------
     >>> import numpy as np
-    >>> backend = MatchingBackend(
-        array_backend = np,
-        default_dtype = np.float32,
-        complex_dtype = np.complex64,
-        default_dtype_int = np.int32
-    )
+    >>> from tme.backends import NumpyFFTWBackend
+    >>> backend = NumpyFFTWBackend(
+    >>>     array_backend=np,
+    >>>     float_dtype=np.float32,
+    >>>     complex_dtype=np.complex64,
+    >>>     int_dtype=np.int32
+    >>> )
     >>> arr = backend.array([1, 2, 3])
     >>> print(arr)
     [1 2 3]
 
     Notes
     -----
-        Developers should be aware of potential naming conflicts between methods and
-        attributes of this class and those of the provided backend.
+    Developers should be aware of potential naming conflicts between methods and
+    attributes of this class and those of the provided backend.
     """
 
     def __init__(
         self,
         array_backend,
-        default_dtype: type,
+        float_dtype: type,
         complex_dtype: type,
-        default_dtype_int: type,
+        int_dtype: type,
+        overflow_safe_dtype: type,
     ):
         self._array_backend = array_backend
-        self._default_dtype = default_dtype
+        self._float_dtype = float_dtype
         self._complex_dtype = complex_dtype
-        self._default_dtype_int = default_dtype_int
+        self._int_dtype = int_dtype
+        self._overflow_safe_dtype = overflow_safe_dtype
+
+        self._fundamental_dtypes = {
+            int: self._int_dtype,
+            float: self._float_dtype,
+            complex: self._complex_dtype,
+        }
 
     def __getattr__(self, name: str):
         """
@@ -115,27 +141,8 @@ class MatchingBackend(ABC):
         base_attributes.extend(dir(self._array_backend))
         return sorted(base_attributes)
 
-    @staticmethod
-    def free_sharedarr(link: shared_memory.SharedMemory):
-        """
-        Free shared array at link.
-
-        Parameters
-        ----------
-        link : shared_memory.SharedMemory
-            The shared memory link to be freed.
-        """
-        if type(link) is not shared_memory.SharedMemory:
-            return None
-        try:
-            link.close()
-            link.unlink()
-        # Shared memory has been freed already
-        except FileNotFoundError:
-            pass
-
     @abstractmethod
-    def to_backend_array(self, arr: NDArray) -> ArrayLike:
+    def to_backend_array(self, arr: NDArray) -> BackendArray:
         """
         Convert a numpy array instance to backend array type.
 
@@ -146,7 +153,7 @@ class MatchingBackend(ABC):
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             An array of the specified backend.
 
         See Also
@@ -156,13 +163,13 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def to_numpy_array(self, arr: ArrayLike) -> NDArray:
+    def to_numpy_array(self, arr: BackendArray) -> NDArray:
         """
         Convert an array of given backend to a numpy array.
 
         Parameters
         ----------
-        arr : NDArray
+        arr : BackendArray
             The array instance to be converted.
 
         Returns
@@ -177,19 +184,19 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def to_cpu_array(self, arr: ArrayLike) -> NDArray:
+    def to_cpu_array(self, arr: BackendArray) -> BackendArray:
         """
         Convert an array of a given backend to a CPU array of that backend.
 
         Parameters
         ----------
-        arr : NDArray
+        arr : BackendArray
             The array instance to be converted.
 
         Returns
         -------
-        NDArray
-            The numpy array equivalent of arr.
+        BackendArray
+            The CPU array equivalent of arr.
 
         See Also
         --------
@@ -197,6 +204,22 @@ class MatchingBackend(ABC):
         :py:meth:`MatchingBackend.to_backend_array`
         """
 
+    def get_fundamental_dtype(self, arr: BackendArray) -> type:
+        """
+        Given an array instance, returns the corresponding fundamental python type,
+        i.e., int, float or complex.
+
+        Parameters
+        ----------
+        arr : BackendArray
+            Input data.
+
+        Returns
+        -------
+        type
+            Data type.
+        """
+
     @abstractmethod
     def free_cache(self):
         """
@@ -204,428 +227,457 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def add(self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None) -> ArrayLike:
+    def add(
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
-        Interface for element-wise addition of arrays.
+        Element-wise addition of arrays.
 
         Parameters
         ----------
-        arr1 : ArrayLike
+        arr1 : BackendArray
             Input array.
-        arr2 : ArrayLike
+        arr2 : BackendArray
             Input array.
-        out : ArrayLike, optional
-            Optional output array to store the result. If provided, it must have a shape
-            that the inputs broadcast to.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Element-wise sum of the input arrays.
         """
 
     @abstractmethod
     def subtract(
-        self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None
-    ) -> ArrayLike:
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
-        Interface for element-wise subtraction of arrays.
+        Element-wise subtraction of arrays.
 
         Parameters
         ----------
-        arr1 : ArrayLike
+        arr1 : BackendArray
             The minuend array.
-        arr2 : ArrayLike
+        arr2 : BackendArray
             The subtrahend array.
-        out : ArrayLike, optional
-            Optional output array to store the result. If provided, it must have a shape
-            that the inputs broadcast to.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Element-wise difference of the input arrays.
         """
 
     @abstractmethod
     def multiply(
-        self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None
-    ) -> ArrayLike:
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
-        Interface for element-wise multiplication of arrays.
+        Element-wise multiplication of arrays.
 
         Parameters
         ----------
-        arr1 : ArrayLike
+        arr1 : BackendArray
             Input array.
-        arr2 : ArrayLike
+        arr2 : BackendArray
             Input array.
-        out : ArrayLike, optional
-            Optional output array to store the result. If provided, it must have a shape
-            that the inputs broadcast to.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Element-wise product of the input arrays.
         """
 
     @abstractmethod
     def divide(
-        self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None
-    ) -> ArrayLike:
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
-        Interface for element-wise division of arrays.
+        Element-wise division of arrays.
 
         Parameters
         ----------
-        arr1 : ArrayLike
-            The numerator array.
-        arr2 : ArrayLike
-            The denominator array.
-        out : ArrayLike, optional
-            Optional output array to store the result. If provided, it must have a shape
-            that the inputs broadcast to.
+        arr1 : BackendArray
+            The dividend array.
+        arr2 : BackendArray
+            The divisor array.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Element-wise quotient of the input arrays.
         """
 
     @abstractmethod
-    def mod(self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None) -> ArrayLike:
+    def mod(
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
-        Interface for element-wise modulus of arrays.
+        Element-wise modulus of arrays.
 
         Parameters
         ----------
-        arr1 : ArrayLike
-            The numerator array.
-        arr2 : ArrayLike
-            The denominator array.
-        out : ArrayLike, optional
-            Optional output array to store the result. If provided, it must have a shape
-            that the inputs broadcast to.
+        arr1 : BackendArray
+            The dividend array.
+        arr2 : BackendArray
+            The divisor array.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Element-wise modulus of the input arrays.
         """
 
     @abstractmethod
-    def sum(self, arr: ArrayLike) -> Scalar:
+    def einsum(
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
-        Compute the sum of array elements.
+        Compute the einstein notation based summation.
 
         Parameters
         ----------
-        arr : ArrayLike
-            The array whose sum should be computed.
+        subscripts : str
+            Specifies the subscripts for summation (see  :obj:`numpy.einsum`).
+        arr1, arr2 : BackendArray
+            Input data.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        Scalar
-            Sum of the arr.
+        BackendArray
+            Einsum of input arrays.
         """
 
     @abstractmethod
-    def einsum(
-        self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None
-    ) -> ArrayLike:
+    def sum(
+        self, arr: BackendArray, axis: Tuple[int] = None
+    ) -> Union[BackendArray, Scalar]:
         """
-        Interface for einstein notation based summation.
+        Compute the sum of array elements.
 
         Parameters
         ----------
-        subscripts : str
-            Specifies the subscripts for summation as comma separated
-            list of subscript label
-        arr1 : ArrayLike
-            Input array.
-        arr2 : ArrayLike
-            Input array.
-        out : ArrayLike, optional
-            Optional output array to store the result. If provided, it must have a shape
-            that the inputs broadcast to.
+        arr : BackendArray
+            Input data.
+        axis : int or tuple of ints, optional
+            Axis or axes to perform the operation on. Default is all.
 
         Returns
         -------
-        ArrayLike
-            Element-wise sum of the input arrays.
+        Union[BackendArray, Scalar]
+            Sum of ``arr``.
         """
 
     @abstractmethod
-    def mean(self, arr: ArrayLike) -> Scalar:
+    def mean(
+        self, arr: BackendArray, axis: Tuple[int] = None
+    ) -> Union[BackendArray, Scalar]:
         """
         Compute the mean of array elements.
 
         Parameters
         ----------
-        arr : ArrayLike
-            The array whose mean should be computed.
+        arr : BackendArray
+            Input data.
+        axis : int or tuple of ints, optional
+            Axis or axes to perform the operation on. Default is all.
 
         Returns
         -------
-        Scalar
-            Mean value of the arr.
+        Union[BackendArray, Scalar]
+            Mean of ``arr``.
         """
 
     @abstractmethod
-    def std(self, arr: ArrayLike, axis: Scalar) -> Scalar:
+    def std(
+        self, arr: BackendArray, axis: Tuple[int] = None
+    ) -> Union[BackendArray, Scalar]:
         """
         Compute the standad deviation of array elements.
 
         Parameters
         ----------
-        arr : Scalar
-            The array whose standard deviation should be computed.
-        axis : Scalar
-            Axis to perform the operation on.
+        arr : BackendArray
+            Input data.
+        axis : int or tuple of ints, optional
+            Axis or axes to perform the operation on. Default is all.
 
         Returns
         -------
-        Scalar
-            The standard deviation of arr.
+        Union[BackendArray, Scalar]
+            Standard deviation of ``arr``.
         """
 
     @abstractmethod
-    def max(self, arr: ArrayLike) -> Scalar:
+    def max(
+        self, arr: BackendArray, axis: Tuple[int] = None
+    ) -> Union[BackendArray, Scalar]:
         """
         Compute the maximum of array elements.
 
         Parameters
         ----------
-        arr : Scalar
-            The array whose maximum should be computed.
+        arr : BackendArray
+            Input data.
+        axis : int or tuple of ints, optional
+            Axis or axes to perform the operation on. Default is all.
 
         Returns
         -------
-        Scalar
-            The maximum of arr.
+        Union[BackendArray, Scalar]
+            Maximum of ``arr``.
         """
 
     @abstractmethod
-    def min(self, arr: ArrayLike) -> Scalar:
+    def min(
+        self, arr: BackendArray, axis: Tuple[int] = None
+    ) -> Union[BackendArray, Scalar]:
         """
         Compute the minimum of array elements.
 
         Parameters
         ----------
-        arr : Scalar
-            The array whose maximum should be computed.
+        arr : BackendArray
+            Input data.
+        axis : int or tuple of ints, optional
+            Axis or axes to perform the operation on. Default is all.
 
         Returns
         -------
-        Scalar
-            The maximum of arr.
+        Union[BackendArray, Scalar]
+            Minimum of ``arr``.
         """
 
     @abstractmethod
     def maximum(
-        self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None
-    ) -> Scalar:
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
         Compute the element wise maximum of arr1 and arr2.
 
         Parameters
         ----------
-        arr1, arr2 : ArrayLike
-            Arrays for which element wise maximum will be computed.
-        out : ArrayLike, optional
-            Output array.
+        arr1, arr2 : BackendArray
+            Input data.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
-            Element wise maximum of arr1 and arr2.
+        BackendArray
+            Element wise maximum of ``arr1`` and ``arr2``.
         """
 
     @abstractmethod
     def minimum(
-        self, arr1: ArrayLike, arr2: ArrayLike, out: ArrayLike = None
-    ) -> Scalar:
+        self, arr1: BackendArray, arr2: BackendArray, out: BackendArray = None
+    ) -> BackendArray:
         """
         Compute the element wise minimum of arr1 and arr2.
 
         Parameters
         ----------
-        arr1, arr2 : ArrayLike
-            Arrays for which element wise minimum will be computed.
-        out : ArrayLike, optional
-            Output array.
+        arr1, arr2 : BackendArray
+            Input data.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Element wise minimum of arr1 and arr2.
         """
 
     @abstractmethod
-    def sqrt(self, arr: ArrayLike) -> ArrayLike:
+    def sqrt(self, arr: BackendArray, out: BackendArray = None) -> BackendArray:
         """
         Compute the square root of array elements.
 
         Parameters
         ----------
-        arr : ArrayLike
-            The array whose square root should be computed.
+        arr : BackendArray
+            Input data.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
-            The squared root of the array.
+        BackendArray
+            Square root of ``arr``.
         """
 
     @abstractmethod
-    def square(self, arr: ArrayLike) -> ArrayLike:
+    def square(self, arr: BackendArray, out: BackendArray = None) -> BackendArray:
         """
         Compute the square of array elements.
 
         Parameters
         ----------
-        arr : ArrayLike
-            The array that should be squared.
+        arr : BackendArray
+            Input data.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
-            The squared arr.
+        BackendArray
+            Square of ``arr``.
         """
 
     @abstractmethod
-    def abs(self, arr: ArrayLike) -> ArrayLike:
+    def abs(self, arr: BackendArray, out: BackendArray = None) -> BackendArray:
         """
         Compute the absolute of array elements.
 
         Parameters
         ----------
-        arr : ArrayLike
-            The array whose absolte should be computed.
+        arr : BackendArray
+            Input data.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
-            The absolute of the array.
+        BackendArray
+            Absolute value of ``arr``.
         """
 
     @abstractmethod
-    def transpose(self, arr: ArrayLike) -> ArrayLike:
+    def transpose(self, arr: BackendArray) -> BackendArray:
         """
         Compute the transpose of arr.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
+        arr : BackendArray
+            Input data.
 
         Returns
         -------
-        ArrayLike
-            The transpose of arr.
+        BackendArray
+            Transpose of ``arr``.
         """
 
-    def power(self, *args, **kwargs) -> ArrayLike:
+    def power(
+        self,
+        arr: BackendArray = None,
+        power: BackendArray = None,
+        out: BackendArray = None,
+        *args,
+        **kwargs,
+    ) -> BackendArray:
         """
         Compute the n-th power of an array.
 
+        Parameters
+        ----------
+        arr : BackendArray
+            Input data.
+        power : BackendArray
+            Power to raise ``arr`` to.
+        arr : BackendArray
+            Output array to write the result to. Returns a new array by default.
+
         Returns
         -------
-        ArrayLike
-            The n-th power of the array
+        BackendArray
+            N-th power of ``arr``.
         """
 
-    def tobytes(self, arr: ArrayLike) -> str:
+    def tobytes(self, arr: BackendArray) -> str:
         """
         Compute the bytestring representation of arr.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
+        arr : BackendArray
+            Input data.
 
         Returns
         -------
         str
-            Bytestring representation of arr.
+            Bytestring representation of ``arr``.
         """
 
     @abstractmethod
-    def size(self, arr: ArrayLike) -> int:
+    def size(self, arr: BackendArray) -> int:
         """
         Compute the number of elements of arr.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
+        arr : BackendArray
+            Input data.
 
         Returns
         -------
         int
-            The number of elements in arr.
+            Number of elements in ``arr``.
         """
 
     @abstractmethod
-    def fill(self, arr: ArrayLike, value: float) -> ArrayLike:
+    def fill(self, arr: BackendArray, value: Scalar) -> None:
         """
-        Fill arr with value.
+        Fills ``arr`` in-place with a given value.
 
         Parameters
         ----------
-        arr : ArrayLike
-            The array that should be filled.
-        value : float
-            The value with which to fill the array.
-
-        Returns
-        -------
-        ArrayLike
-            The array filled with the given value.
+        arr : BackendArray
+            Input data.
+        value : Scalar
+            The value to fill the array with.
         """
 
     @abstractmethod
-    def zeros(self, shape: Tuple[int], dtype: type) -> ArrayLike:
+    def zeros(self, shape: Tuple[int], dtype: type) -> BackendArray:
         """
         Returns an aligned array of zeros with specified shape and dtype.
 
         Parameters
         ----------
-        shape : Tuple[int]
+        shape : tuple of ints.
             Desired shape for the array.
         dtype : type
             Desired data type for the array.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Byte-aligned array of zeros with specified shape and dtype.
         """
 
     @abstractmethod
-    def full(self, shape: Tuple[int], dtype: type, fill_value: Scalar) -> ArrayLike:
+    def full(self, shape: Tuple[int], dtype: type, fill_value: Scalar) -> BackendArray:
         """
-        Returns an aligned array of zeros with specified shape and dtype.
+        Returns an array filled with fill_value of specified shape and dtype.
 
         Parameters
         ----------
-        shape : Tuple[int]
+        shape : tuple of ints.
             Desired shape for the array.
         dtype : type
             Desired data type for the array.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Byte-aligned array of zeros with specified shape and dtype.
         """
 
     @abstractmethod
     def eps(self, dtype: type) -> Scalar:
         """
-        Returns the eps defined as diffeerence between 1.0 and the next
-        representable floating point value larger than 1.0.
+        Returns the minimal difference representable by dtype.
 
         Parameters
         ----------
@@ -635,7 +687,7 @@ class MatchingBackend(ABC):
         Returns
         -------
         Scalar
-            The eps for the given data type
+            The eps for the given data type.
         """
 
     @abstractmethod
@@ -646,325 +698,311 @@ class MatchingBackend(ABC):
         Parameters
         ----------
         dtype : type
-            Datatype for which the number of bytes is to be determined.
-            This is typically a data type like `np.float32` or `np.int64`.
+            Data type to determine the bytesize of.
 
         Returns
         -------
         int
             Number of bytes occupied by the datatype.
-
-        Examples
-        --------
-        >>> MatchingBackend.datatype_bytes(np.float32)
-        4
         """
 
     @abstractmethod
     def clip(
-        self, arr: ArrayLike, a_min: Scalar, a_max: Scalar, out: ArrayLike = None
-    ) -> ArrayLike:
+        self, arr: BackendArray, a_min: Scalar, a_max: Scalar, out: BackendArray = None
+    ) -> BackendArray:
         """
         Clip elements of arr.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
+        arr : BackendArray
+            Input data.
         a_min : Scalar
             Lower bound.
         a_max : Scalar
             Upper bound.
-        out : ArrayLike, optional
-            Output array.
+        out : BackendArray, optional
+            Output array to write the result to. Returns a new array by default.
 
         Returns
         -------
-        ArrayLike
-            Clipped arr.
+        BackendArray
+            Clipped ``arr``.
         """
 
     @abstractmethod
-    def flip(
-        self, arr: ArrayLike, a_min: Scalar, a_max: Scalar, out: ArrayLike = None
-    ) -> ArrayLike:
-        """
-        Flip the elements of arr.
-
-        Parameters
-        ----------
-        arr : ArrayLike
-            Input array.
-
-        Returns
-        -------
-        ArrayLike
-            Flipped version of arr.
-        """
-
-    @abstractmethod
-    def astype(arr: ArrayLike, dtype: type) -> ArrayLike:
+    def astype(arr: BackendArray, dtype: type) -> BackendArray:
         """
         Change the datatype of arr.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
+        arr : BackendArray
+            Input data.
         dtype : type
             Target data type.
 
         Returns
         -------
-        ArrayLike
-            Flipped version of arr.
+        BackendArray
+            Freshly allocated array containing the data of ``arr`` in ``dtype``.
         """
 
     @abstractmethod
-    def arange(self, *args, **kwargs) -> ArrayLike:
+    def arange(
+        self, stop: Scalar, start: Scalar = 0, step: Scalar = 1, *args, **kwargs
+    ) -> BackendArray:
         """
-        Arange values in increasing order.
+        Arange values in evenly spaced interval.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
-        dtype : type
-            Target data type.
+        stop : Scalar
+            End of the interval.
+        start : Scalar
+            Start of the interval, zero by default.
+        step : Scalar
+            Interval step size, one by default.
 
         Returns
         -------
-        ArrayLike
-            Flipped version of arr.
+        BackendArray
+            Array of evenly spaced values in specified interval.
         """
 
-    def stack(self, *args, **kwargs) -> ArrayLike:
+    def stack(self, *args, **kwargs) -> BackendArray:
         """
-        Join a sequence of ArrayLike objects along a specified axis.
+        Join a sequence of objects along a new axis.
+
+        Parameters
+        ----------
+        arr : BackendArray
+            Sequence of arrays.
+        axis : int, optional
+            Axis along which to stack the input arrays.
 
         Returns
         -------
-        ArrayLike
-            The joined input objects.
+        BackendArray
+            Stacked input data.
         """
 
     @abstractmethod
-    def concatenate(self, *args, **kwargs) -> ArrayLike:
+    def concatenate(self, *args, **kwargs) -> BackendArray:
         """
-        Concatenates arrays along axis.
+        Join a sequence of objects along an existing axis.
 
         Parameters
         ----------
-        arr * : ArrayLike
-            Input arrays
+        arr : BackendArray
+            Sequence of arrays.
+        axis : int
+            Axis along which to stack the input arrays.
 
         Returns
         -------
-        ArrayLike
-            Concatenated input arrays
+        BackendArray
+            Concatenated input data.
         """
 
     @abstractmethod
-    def repeat(self, *args, **kwargs) -> ArrayLike:
+    def repeat(self, *args, **kwargs) -> BackendArray:
         """
-        Repeat each array elements after themselves.
+        Repeat each array element a specified number of times.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array
+        arr : BackendArray
+            Input data.
+        repeats : int or tuple of ints
+            Number of each repetitions along axis.
 
         Returns
         -------
-        ArrayLike
-            Repeated input array
+        BackendArray
+            Repeated ``arr``.
         """
 
     @abstractmethod
-    def topk_indices(self, arr: NDArray, k: int) -> ArrayLike:
+    def topk_indices(self, arr: NDArray, k: int) -> BackendArray:
         """
-        Compute indices of largest elements.
+        Determinces the indices of largest elements.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
-        dtype : type
-            Target data type.
+        arr : BackendArray
+            Input data.
+        k : int
+            Number of maxima to determine.
 
         Returns
         -------
-        ArrayLike
-            Flipped version of arr.
+        BackendArray
+            Indices of ``k`` largest elements in ``arr``.
         """
 
-    def indices(self, *args, **kwargs) -> ArrayLike:
+    def indices(self, *args, **kwargs) -> BackendArray:
         """
         Creates an array representing the index grid of an input.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             The index grid.
         """
 
     @abstractmethod
-    def roll(self, *args, **kwargs) -> ArrayLike:
+    def roll(self, *args, **kwargs) -> BackendArray:
         """
         Roll array elements along a specified axis.
 
         Parameters
         ----------
-        *args : tuple
-            Generic arguments.
-        **kwargs : dict
-            Generic keyword arguments.
+        a : BackendArray
+            Input data.
+        shift : int or tuple of ints, optional
+            Shift along each axis.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Array with elements rolled.
+        """
 
-        See Also
-        --------
-        numpy.roll : For more detailed documentation on arguments and behavior.
+    @abstractmethod
+    def where(condition, *args) -> BackendArray:
+        """
+        Return elements from input depending on ``condition``.
 
-        Examples
-        --------
-        >>> import numpy as np
-        >>> x = np.array([1, 2, 3, 4, 5])
-        >>> np.roll(x, 2)
-        array([4, 5, 1, 2, 3])
+        Parameters
+        ----------
+        condition : BackendArray
+            Binary condition array.
+        *args : BackendArray
+            Values to choose from.
+
+        Returns
+        -------
+        BackendArray
+            Array of elements according to ``condition``.
         """
 
     @abstractmethod
-    def unique(self, *args, **kwargs) -> Tuple[ArrayLike]:
+    def unique(
+        self,
+        arr: BackendArray,
+        return_index: bool = False,
+        return_inverse: bool = False,
+        return_counts: bool = False,
+        axis: Tuple[int] = None,
+        *args,
+        **kwargs,
+    ) -> Tuple[BackendArray]:
         """
         Find the unique elements of an array.
 
         Parameters
         ----------
-        *args : tuple
-            Generic arguments.
-        **kwargs : dict
-            Generic keyword arguments.
+        arr : BackendArray
+            Input data.
+        return_index : bool, optional
+            Return indices that resulted in unique array, False by default.
+        return_inverse : bool, optional
+            Return indices to reconstruct the input, False by default.
+        return_counts : bool, optional
+            Return number of occurences of each unique element, False by default.
+        axis : int or tuple of ints, optional
+            Axis or axes to perform the operation on. Default is all.
 
         Returns
         -------
-        ArrayLike or tuple of ArrayLike
+        BackendArray or tuple of BackendArray
             If `return_index`, `return_inverse`, and `return_counts` keyword
-            arguments are all False (the default), this will be an ArrayLike object
+            arguments are all False (the default), this will be an BackendArray object
             of the sorted unique values. Otherwise, it's a tuple with one
             or more arrays as specified by those keyword arguments.
-
-        See Also
-        --------
-        numpy.unique : For more detailed documentation on arguments and behavior.
-
-        Examples
-        --------
-        >>> import numpy as np
-        >>> x = np.array([1, 2, 3, 2, 3, 4])
-        >>> np.unique(x)
-        array([1, 2, 3, 4])
         """
 
     @abstractmethod
-    def argsort(self, *args, **kwargs) -> ArrayLike:
+    def argsort(self, *args, **kwargs) -> BackendArray:
         """
-        Perform argsort of arr.
+        Compute the indices to sort a given input array.
 
         Parameters
         ----------
-        arr : ArrayLike
+        arr : BackendArray
             Input array.
         dtype : type
             Target data type.
 
         Returns
         -------
-        ArrayLike
-            Flipped version of arr.
+        BackendArray
+            Indices that would sort the input data.
         """
 
     @abstractmethod
-    def unravel_index(self, indices: ArrayLike, shape: Tuple[int]) -> Tuple[ArrayLike]:
+    def unravel_index(
+        self, indices: BackendArray, shape: Tuple[int]
+    ) -> Tuple[BackendArray]:
         """
         Convert flat index to array indices.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
-        dtype : type
-            Target data type.
+        indices : BackendArray
+            Input data.
+        shape : tuple of ints
+            Shape of the array used for unraveling.
 
         Returns
         -------
-        ArrayLike
-            Flipped version of arr.
+        BackendArray
+            Array indices.
         """
 
     @abstractmethod
-    def tril_indices(self, *args, **kwargs) -> ArrayLike:
+    def tril_indices(self, *args, **kwargs) -> BackendArray:
         """
         Compute indices of upper triangular matrix
 
         Parameters
         ----------
-        arr : ArrayLike
+        arr : BackendArray
             Input array.
         dtype : type
             Target data type.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Flipped version of arr.
         """
 
     @abstractmethod
     def max_filter_coordinates(
-        self, score_space: ArrayLike, min_distance: Tuple[int]
-    ) -> ArrayLike:
+        self, score_space: BackendArray, min_distance: Tuple[int]
+    ) -> BackendArray:
         """
         Identifies local maxima in score_space separated by min_distance.
 
         Parameters
         ----------
-        score_space : ArrayLike
+        score_space : BackendArray
             Input score space.
         min_distance : tuple of ints
             Minimum distance along each array axis.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Identified local maxima.
         """
 
     @abstractmethod
-    def preallocate_array(self, shape: Tuple[int], dtype: type) -> ArrayLike:
-        """
-        Returns an aligned array of zeros with specified shape and dtype.
-
-        Parameters
-        ----------
-        shape : Tuple[int]
-            Desired shape for the array.
-        dtype : type
-            Desired data type for the array.
-
-        Returns
-        -------
-        ArrayLike
-            Byte-aligned array of zeros with specified shape and dtype.
-        """
-
-    @abstractmethod
-    def sharedarr_to_arr(
+    def from_sharedarr(
         self, shape: Tuple[int], dtype: str, shm: shared_memory.SharedMemory
-    ) -> ArrayLike:
+    ) -> BackendArray:
         """
         Returns an array of given shape and dtype from shared memory location.
 
@@ -979,41 +1017,41 @@ class MatchingBackend(ABC):
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Array of the specified shape and dtype from the shared memory location.
         """
 
     @abstractmethod
-    def arr_to_sharedarr(
-        self, arr: type, shared_memory_handler: type = None
-    ) -> shared_memory.SharedMemory:
+    def to_sharedarr(self, arr: type, shared_memory_handler: type = None) -> shm_type:
         """
-        Converts an array to an object shared in memory.
+        Converts an array to an object shared in memory. The output of this function
+        will only be passed to :py:meth:`from_sharedarr`, hence the return values can
+        be modified in particular backends to match the expected input data.
 
         Parameters
         ----------
-        arr : ArrayLike
+        arr : BackendArray
             Numpy array to convert.
         shared_memory_handler : type, optional
             The type of shared memory handler. Default is None.
 
         Returns
         -------
-        shared_memory.SharedMemory
-            The shared memory object containing the numpy array.
+        Tuple[shared_memory.SharedMemory, tuple of ints, dtype]
+            The shared memory object containing the numpy array, its shape and dtype.
         """
 
     @abstractmethod
     def topleft_pad(
-        self, arr: ArrayLike, shape: Tuple[int], padval: int = 0
-    ) -> ArrayLike:
+        self, arr: BackendArray, shape: Tuple[int], padval: int = 0
+    ) -> BackendArray:
         """
         Returns an array that has been padded to a specified shape with a padding
         value at the top-left corner.
 
         Parameters
         ----------
-        arr : ArrayLike
+        arr : BackendArray
             Input array to be padded.
         shape : Tuple[int]
             Desired shape for the output array.
@@ -1022,7 +1060,7 @@ class MatchingBackend(ABC):
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Array that has been padded to the specified shape.
         """
 
@@ -1036,7 +1074,11 @@ class MatchingBackend(ABC):
         **kwargs,
     ) -> Tuple[Callable, Callable]:
         """
-        Build forward and inverse real fourier transform functions.
+        Build forward and inverse real fourier transform functions. The returned
+        callables have two parameters ``arr`` and ``out`` which correspond to the
+        input and output of the Fourier transform. The methods return the output
+        of the respective function call, regardless of ``out`` being provided or not,
+        analogous to most numpy functions.
 
         Parameters
         ----------
@@ -1050,35 +1092,36 @@ class MatchingBackend(ABC):
             Numpy dtype of the inverse fourier transform.
         complex_dtype : dtype
             Numpy dtype of the fourier transform.
+        inverse_fast_shape : tuple, optional
+            Output shape of the inverse Fourier transform. By default fast_shape.
         fftargs : dict, optional
             Dictionary passed to pyFFTW builders.
-        temp_real : ArrayLike, optional
+        temp_real : NDArray, optional
             Temporary real numpy array, by default None.
-        temp_fft : ArrayLike, optional
+        temp_fft : NDArray, optional
             Temporary fft numpy array, by default None.
 
         Returns
         -------
         tuple
-            Tuple containing function pointers for forward and inverse real
-            fourier transform
+            Tuple of callables for forward and inverse real Fourier transform.
         """
 
-    def extract_center(self, arr: ArrayLike, newshape: Tuple[int]) -> ArrayLike:
+    def extract_center(self, arr: BackendArray, newshape: Tuple[int]) -> BackendArray:
         """
         Extract the centered portion of an array based on a new shape.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
+        arr : BackendArray
+            Input data.
         newshape : tuple
             Desired shape for the central portion.
 
         Returns
         -------
-        ArrayLike
-            Central portion of the array with shape `newshape`.
+        BackendArray
+            Central portion of the array with shape ``newshape``.
         """
 
     @abstractmethod
@@ -1104,19 +1147,42 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def rotate_array(self, *args, **kwargs):
+    def rigid_transform(
+        self,
+        arr: BackendArray,
+        rotation_matrix: BackendArray,
+        arr_mask: Optional[BackendArray] = None,
+        translation: Optional[BackendArray] = None,
+        use_geometric_center: bool = True,
+        out: Optional[BackendArray] = None,
+        out_mask: Optional[BackendArray] = None,
+        order: int = 3,
+        **kwargs,
+    ) -> Tuple[BackendArray, Optional[BackendArray]]:
         """
-        Perform a rigid transform of arr.
+        Performs a rigid transformation.
 
         Parameters
         ----------
-        arr : ArrayLike
-            Input array.
-
-        Returns
-        -------
-        ArrayLike
-            Transformed version of arr.
+        arr : BackendArray
+            The input array to be rotated.
+        arr_mask : BackendArray, optional
+            The mask of `arr` that will be equivalently rotated.
+        rotation_matrix : BackendArray
+            The rotation matrix to apply (d, d).
+        translation : BackendArray
+            The translation to apply (d,).
+        use_geometric_center : bool, optional
+            Whether rotation should be performed over the center of mass.
+        out : BackendArray, optional
+            Location into which the rotation of ``arr`` is written.
+        out_mask : BackendArray, optional
+            Location into which the rotation of ``arr_mask`` is written.
+        order : int, optional
+            Interpolation order, one is linear and three is cubic. Specific
+            meaning depends on backend.
+        kwargs : dict, optional
+            Keyword arguments relevant to particular backend implementations.
         """
 
     @abstractmethod
@@ -1128,29 +1194,39 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def reverse(arr: ArrayLike) -> ArrayLike:
+    def reverse(arr: BackendArray) -> BackendArray:
         """
         Reverse the order of elements in an array along all its axes.
 
         Parameters
         ----------
-        arr : ArrayLike
+        arr : BackendArray
             Input array.
 
         Returns
         -------
-        ArrayLike
+        BackendArray
             Reversed array.
         """
 
     @abstractmethod
-    def set_device(device_index: int):
+    def set_device(device_index: int) -> Generator:
         """
-        Set the active device context and operate as context manager for it
+        Context manager that sets active compute device device for operations.
+
+        Parameters
+        ----------
+        device_index : int
+            Index of the device to be set as active.
         """
 
     @abstractmethod
     def device_count() -> int:
         """
-        Return the number of available compute devices.
+        Return the number of available compute devices considered by the backend.
+
+        Returns
+        -------
+        int
+            Number of available devices.
         """