Add overloads. Fix typings (#16)

* Add overloads. Fix examples. Add missing reqs

Author: roaffix

Makefile

@@ -14,6 +14,10 @@ endif
 version : 
 	@python -c 'from arrayfire.version import VERSION; print(f"ArrayFire Python v{VERSION}")'
 
+.PHONY : build
+build :
+	@python -m build
+
 # Dev
 
 .PHONY : pre-commit

arrayfire/__init__.py

@@ -562,13 +562,23 @@
     "approx1_uniform",
     "approx2",
     "approx2_uniform",
+    "convolve1",
+    "convolve2",
+    "convolve2_nn",
+    "convolve2_separable",
+    "convolve3",
 ]
 
 from arrayfire.library.signal_processing import (
     approx1,
     approx1_uniform,
     approx2,
     approx2_uniform,
+    convolve1,
+    convolve2,
+    convolve2_nn,
+    convolve2_separable,
+    convolve3,
     fft,
     fft2,
     fft2_c2r,

arrayfire/array_api/_sorting_functions.py

@@ -35,7 +35,7 @@ def argsort(x: Array, /, *, axis: int = -1, descending: bool = False, stable: bo
     if axis == -1:
         axis = x.ndim - 1
 
-    _, indices = af.sort(x._array, axis=axis, is_ascending=not descending, is_index_array=True)  # type: ignore[misc]
+    _, indices = af.sort(x._array, axis=axis, is_ascending=not descending, is_index_array=True)
     return Array._new(indices)
 
 

arrayfire/array_object.py

@@ -549,41 +549,41 @@ def __ne__(self, other: int | float | bool | Array, /) -> Array:  # type: ignore
 
     # Reflected Arithmetic Operators
 
-    def __radd__(self, other: Array, /) -> Array:
+    def __radd__(self, other: int | float | Array, /) -> Array:
         """
         Return other + self.
         """
         return process_c_function(other, self, wrapper.add)
 
-    def __rsub__(self, other: Array, /) -> Array:
+    def __rsub__(self, other: int | float | Array, /) -> Array:
         """
         Return other - self.
         """
         return process_c_function(other, self, wrapper.sub)
 
-    def __rmul__(self, other: Array, /) -> Array:
+    def __rmul__(self, other: int | float | Array, /) -> Array:
         """
         Return other * self.
         """
         return process_c_function(other, self, wrapper.mul)
 
-    def __rtruediv__(self, other: Array, /) -> Array:
+    def __rtruediv__(self, other: int | float | Array, /) -> Array:
         """
         Return other / self.
         """
         return process_c_function(other, self, wrapper.div)
 
-    def __rfloordiv__(self, other: Array, /) -> Array:
+    def __rfloordiv__(self, other: int | float | Array, /) -> Array:
         # TODO
         return NotImplemented
 
-    def __rmod__(self, other: Array, /) -> Array:
+    def __rmod__(self, other: int | float | Array, /) -> Array:
         """
         Return other % self.
         """
         return process_c_function(other, self, wrapper.mod)
 
-    def __rpow__(self, other: Array, /) -> Array:
+    def __rpow__(self, other: int | float | Array, /) -> Array:
         """
         Return other ** self.
         """
@@ -597,31 +597,31 @@ def __rmatmul__(self, other: Array, /) -> Array:
 
     # Reflected Bitwise Operators
 
-    def __rand__(self, other: Array, /) -> Array:
+    def __rand__(self, other: int | bool | Array, /) -> Array:
         """
         Return other & self.
         """
         return process_c_function(other, self, wrapper.bitand)
 
-    def __ror__(self, other: Array, /) -> Array:
+    def __ror__(self, other: int | bool | Array, /) -> Array:
         """
         Return other | self.
         """
         return process_c_function(other, self, wrapper.bitor)
 
-    def __rxor__(self, other: Array, /) -> Array:
+    def __rxor__(self, other: int | bool | Array, /) -> Array:
         """
         Return other ^ self.
         """
         return process_c_function(other, self, wrapper.bitxor)
 
-    def __rlshift__(self, other: Array, /) -> Array:
+    def __rlshift__(self, other: int | Array, /) -> Array:
         """
         Return other << self.
         """
         return process_c_function(other, self, wrapper.bitshiftl)
 
-    def __rrshift__(self, other: Array, /) -> Array:
+    def __rrshift__(self, other: int | Array, /) -> Array:
         """
         Return other >> self.
         """
@@ -1126,11 +1126,11 @@ def process_c_function(lhs: int | float | Array, rhs: int | float | Array, c_fun
         lhs_array = lhs.arr
         rhs_array = rhs.arr
 
-    elif isinstance(lhs, Array) and isinstance(rhs, (int, float)):
+    elif isinstance(lhs, Array) and isinstance(rhs, int | float):
         lhs_array = lhs.arr
         rhs_array = wrapper.create_constant_array(rhs, lhs.shape, lhs.dtype)
 
-    elif isinstance(lhs, (int, float)) and isinstance(rhs, Array):
+    elif isinstance(lhs, int | float) and isinstance(rhs, Array):
         lhs_array = wrapper.create_constant_array(lhs, rhs.shape, rhs.dtype)
         rhs_array = rhs.arr
 

arrayfire/library/device.py

@@ -47,5 +47,25 @@
     set_device,
     set_kernel_cache_directory,
     set_mem_step_size,
-    sync,
 )
+from arrayfire_wrapper.lib import sync as wrapper_sync
+
+
+def sync(device_id: int | None = None) -> None:
+    """
+    Blocks until all the functions on the specified device have completed execution.
+
+    This function is used to synchronize the program execution with the operations
+    being carried out on a GPU or other computation device, ensuring that all
+    previously submitted operations are complete before the program proceeds.
+
+    Parameters
+    ----------
+    device_id : int | None, optional
+        The ID of the device on which to wait for all operations to complete.
+        If None is provided, the current active device is used. Default is None.
+    """
+    if device_id is None:
+        device_id = get_device()
+
+    wrapper_sync(device_id)

arrayfire/library/linear_algebra.py

@@ -15,7 +15,7 @@
     "solve",
 ]
 
-from typing import cast
+from typing import Literal, cast, overload
 
 import arrayfire_wrapper.lib as wrapper
 from arrayfire_wrapper.lib import is_lapack_available
@@ -24,6 +24,17 @@
 from arrayfire.array_object import afarray_as_array
 from arrayfire.library.constants import MatProp, Norm
 
+# TODO
+# Add missing documentation
+
+
+@overload
+def dot(lhs: Array, rhs: Array, /, *, return_scalar: Literal[True]) -> int | float | complex: ...
+
+
+@overload
+def dot(lhs: Array, rhs: Array, /, *, return_scalar: Literal[False] = False) -> Array: ...
+
 
 def dot(
     lhs: Array,
@@ -184,6 +195,14 @@ def matmul(lhs: Array, rhs: Array, /, lhs_opts: MatProp = MatProp.NONE, rhs_opts
     return cast(Array, wrapper.matmul(lhs.arr, rhs.arr, lhs_opts, rhs_opts))
 
 
+@overload
+def cholesky(array: Array, /, is_upper: bool = True, *, inplace: Literal[True]) -> int: ...
+
+
+@overload
+def cholesky(array: Array, /, is_upper: bool = True, *, inplace: Literal[False] = False) -> tuple[Array, int]: ...
+
+
 def cholesky(array: Array, /, is_upper: bool = True, *, inplace: bool = False) -> int | tuple[Array, int]:
     if inplace:
         return wrapper.cholesky_inplace(array.arr, is_upper)
@@ -192,6 +211,14 @@ def cholesky(array: Array, /, is_upper: bool = True, *, inplace: bool = False) -
     return Array.from_afarray(matrix), info
 
 
+@overload
+def lu(array: Array, /, *, inplace: Literal[True], is_lapack_pivot: bool = True) -> Array: ...
+
+
+@overload
+def lu(array: Array, /, *, inplace: Literal[False] = False, is_lapack_pivot: bool = True) -> tuple[Array, ...]: ...
+
+
 def lu(array: Array, /, *, inplace: bool = False, is_lapack_pivot: bool = True) -> Array | tuple[Array, ...]:
     if inplace:
         return Array.from_afarray(wrapper.lu_inplace(array.arr, is_lapack_pivot))
@@ -200,6 +227,14 @@ def lu(array: Array, /, *, inplace: bool = False, is_lapack_pivot: bool = True)
     return Array.from_afarray(lower), Array.from_afarray(upper), Array.from_afarray(pivot)
 
 
+@overload
+def qr(array: Array, /, *, inplace: Literal[True]) -> Array: ...
+
+
+@overload
+def qr(array: Array, /, *, inplace: Literal[False] = False) -> tuple[Array, ...]: ...
+
+
 def qr(array: Array, /, *, inplace: bool = False) -> Array | tuple[Array, ...]:
     if inplace:
         return Array.from_afarray(wrapper.qr_inplace(array.arr))
@@ -247,5 +282,5 @@ def solve(a: Array, b: Array, /, *, options: MatProp = MatProp.NONE, pivot: None
     return cast(Array, wrapper.solve(a.arr, b.arr, options))
 
 
-# TODO #good_first_issue
-# Add Sparse functions
+# TODO
+# Create issues as #good_first_issue: add Sparse functions

arrayfire/library/signal_processing.py

@@ -21,6 +21,11 @@
     "approx1_uniform",
     "approx2",
     "approx2_uniform",
+    "convolve1",
+    "convolve2",
+    "convolve2_nn",
+    "convolve2_separable",
+    "convolve3",
 ]
 
 from typing import cast

arrayfire/library/statistics.py

@@ -1,13 +1,16 @@
 __all__ = ["corrcoef", "cov", "mean", "median", "stdev", "topk", "var"]
 
-from typing import cast
+from typing import cast, overload
 
 import arrayfire_wrapper.lib as wrapper
 
 from arrayfire import Array
 from arrayfire.array_object import afarray_as_array
 from arrayfire.library.constants import TopK, VarianceBias
 
+# TODO
+# Add missing documentation
+
 
 def corrcoef(x: Array, y: Array, /) -> int | float | complex:
     return wrapper.corrcoef(x.arr, y.arr)
@@ -18,6 +21,22 @@ def cov(x: Array, y: Array, /, *, bias: VarianceBias = VarianceBias.DEFAULT) ->
     return cast(Array, wrapper.cov(x.arr, y.arr, bias))
 
 
+@overload
+def mean(x: Array, /, axis: None = None, *, weights: None = None) -> int | float | complex: ...
+
+
+@overload
+def mean(x: Array, /, axis: int, *, weights: None = None) -> Array: ...
+
+
+@overload
+def mean(x: Array, /, axis: None, *, weights: Array) -> int | float | complex: ...
+
+
+@overload
+def mean(x: Array, /, axis: int, *, weights: Array) -> Array: ...
+
+
 def mean(x: Array, /, axis: None | int = None, *, weights: None | Array = None) -> int | float | complex | Array:
     if weights:
         if axis is None:
@@ -31,13 +50,29 @@ def mean(x: Array, /, axis: None | int = None, *, weights: None | Array = None)
     return Array.from_afarray(wrapper.mean(x.arr, axis))
 
 
+@overload
+def median(x: Array, /, axis: None = None) -> int | float | complex: ...
+
+
+@overload
+def median(x: Array, /, axis: int) -> Array: ...
+
+
 def median(x: Array, /, axis: None | int = None) -> int | float | complex | Array:
     if axis is None:
         return wrapper.median_all(x.arr)
 
     return Array.from_afarray(wrapper.median(x.arr, axis))
 
 
+@overload
+def stdev(x: Array, /, axis: None = None, *, bias: VarianceBias = VarianceBias.DEFAULT) -> int | float | complex: ...
+
+
+@overload
+def stdev(x: Array, /, axis: int, *, bias: VarianceBias = VarianceBias.DEFAULT) -> int | float | complex: ...
+
+
 def stdev(
     x: Array, /, axis: None | int = None, *, bias: VarianceBias = VarianceBias.DEFAULT
 ) -> int | float | complex | Array:
@@ -52,6 +87,26 @@ def topk(x: Array, k: int, /, *, axis: int = 0, order: TopK = TopK.DEFAULT) -> t
     return Array.from_afarray(values), Array.from_afarray(indices)
 
 
+@overload
+def var(
+    x: Array, /, axis: None = None, *, weights: None = None, bias: VarianceBias = VarianceBias.DEFAULT
+) -> int | float | complex: ...
+
+
+@overload
+def var(x: Array, /, axis: int, *, weights: None = None, bias: VarianceBias = VarianceBias.DEFAULT) -> Array: ...
+
+
+@overload
+def var(
+    x: Array, /, axis: None, *, weights: Array, bias: VarianceBias = VarianceBias.DEFAULT
+) -> int | float | complex: ...
+
+
+@overload
+def var(x: Array, /, axis: int, *, weights: Array, bias: VarianceBias = VarianceBias.DEFAULT) -> Array: ...
+
+
 def var(
     x: Array,
     /,