pulled in most recent changes from master

code/approx/approx.c

@@ -36,7 +36,14 @@ STATIC mp_float_t approx_python_call(const mp_obj_type_t *type, mp_obj_t fun, mp
 	return mp_obj_get_float(type->call(fun, nparams+1, 0, fargs));
 }
 
-//| def bisect(fun, a, b, *, xtol=2.4e-7, maxiter=100) -> float:
+//| def bisect(
+//|     fun: Callable[[float], float],
+//|     a: float,
+//|     b: float,
+//|     *,
+//|     xtol: float = 2.4e-7,
+//|     maxiter: int = 100
+//| ) -> float:
 //|     """
 //|     :param callable f: The function to bisect
 //|     :param float a: The left side of the interval
@@ -97,7 +104,14 @@ STATIC mp_obj_t approx_bisect(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(approx_bisect_obj, 3, approx_bisect);
 
-//| def fmin(fun, x0, *, xatol=2.4e-7, fatol=2.4e-7, maxiter=200) -> float:
+//| def fmin(
+//|     fun: Callable[[float], float],
+//|     x0: float,
+//|     *,
+//|     xatol: float = 2.4e-7,
+//|     fatol: float = 2.4e-7,
+//|     maxiter: int = 200
+//| ) -> float:
 //|     """
 //|     :param callable f: The function to bisect
 //|     :param float x0: The initial x value
@@ -317,7 +331,14 @@ MP_DEFINE_CONST_FUN_OBJ_KW(approx_curve_fit_obj, 2, approx_curve_fit);
 
 #endif
 
-//| def interp(x: ulab.array, xp:ulab.array, fp:ulab.array, *, left=None, right=None) -> ulab.array:
+//| def interp(
+//|     x: ulab.array,
+//|     xp: ulab.array,
+//|     fp: ulab.array,
+//|     *,
+//|     left: Optional[float] = None,
+//|     right: Optional[float] = None
+//| ) -> ulab.array:
 //|     """
 //|     :param ulab.array x: The x-coordinates at which to evaluate the interpolated values.
 //|     :param ulab.array xp: The x-coordinates of the data points, must be increasing
@@ -392,7 +413,14 @@ STATIC mp_obj_t approx_interp(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(approx_interp_obj, 2, approx_interp);
 
-//| def newton(fun, x0, *, xtol=2.4e-7, rtol=0.0, maxiter=50) -> float:
+//| def newton(
+//|     fun: Callable[[float], float],
+//|     x0: float,
+//|     *,
+//|     xtol: float = 2.4e-7,
+//|     rtol: float = 0.0,
+//|     maxiter: int = 50
+//| ) -> float:
 //|     """
 //|     :param callable f: The function to bisect
 //|     :param float x0: The initial x value
@@ -447,7 +475,7 @@ STATIC mp_obj_t approx_newton(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(approx_newton_obj, 2, approx_newton);
 
-//| def trapz(y: ulab.array, x=None, dx=1.0) -> float:
+//| def trapz(y: ulab.array, x: Optional[ulab.array] = None, dx: float = 1.0) -> float:
 //|     """
 //|     :param 1D ulab.array y: the values of the dependent variable
 //|     :param 1D ulab.array x: optional, the coordinates of the independent variable. Defaults to uniformly spaced values.

code/compare/compare.c

@@ -134,17 +134,22 @@ static mp_obj_t compare_equal_helper(mp_obj_t x1, mp_obj_t x2, uint8_t comptype)
 
 }
 
-//| def clip(x1, x2, x3):
+//| def clip(
+//|     x1: Union[ulab.array, float],
+//|     x2: Union[ulab.array, float],
+//|     x3: Union[ulab.array, float],
+//| ) -> ulab.array:
 //|     """
 //|     Constrain the values from ``x1`` to be between ``x2`` and ``x3``.
 //|     ``x2`` is assumed to be less than or equal to ``x3``.
+//|
 //|     Arguments may be ulab arrays or numbers.  All array arguments
 //|     must be the same size.  If the inputs are all scalars, a 1-element
 //|     array is returned.
+//|
 //|     Shorthand for ``ulab.maximum(x2, ulab.minimum(x1, x3))``"""
 //|     ...
-//| 
+//|
-
 static mp_obj_t compare_clip(mp_obj_t x1, mp_obj_t x2, mp_obj_t x3) {
 	// Note: this function could be made faster by implementing a single-loop comparison in 
 	// RUN_COMPARE_LOOP. However, that would add around 2 kB of compile size, while we 
@@ -155,37 +160,35 @@ static mp_obj_t compare_clip(mp_obj_t x1, mp_obj_t x2, mp_obj_t x3) {
 
 MP_DEFINE_CONST_FUN_OBJ_3(compare_clip_obj, compare_clip);
 
-//| def equal(x1, x2):
+//| def equal(x1: Union[ulab.array, float], x2: Union[ulab.array, float]) -> List[bool]:
-//|    """Return an array of bool which is true where x1[i] == x2[i] and false elsewhere"""
+//|     """Return an array of bool which is true where x1[i] == x2[i] and false elsewhere"""
-//|    ...
+//|     ...
 //|
-
 static mp_obj_t compare_equal(mp_obj_t x1, mp_obj_t x2) {
 	return compare_equal_helper(x1, x2, MP_BINARY_OP_EQUAL);
 }
 
 MP_DEFINE_CONST_FUN_OBJ_2(compare_equal_obj, compare_equal);
 
-//| def not_equal(x1, x2):
+//| def not_equal(x1: Union[ulab.array, float], x2: Union[ulab.array, float]) -> List[bool]:
-//|    """Return an array of bool which is false where x1[i] == x2[i] and true elsewhere"""
+//|     """Return an array of bool which is false where x1[i] == x2[i] and true elsewhere"""
-//|    ...
+//|     ...
 //|
-
 static mp_obj_t compare_not_equal(mp_obj_t x1, mp_obj_t x2) {
 	return compare_equal_helper(x1, x2, MP_BINARY_OP_NOT_EQUAL);
 }
 
 MP_DEFINE_CONST_FUN_OBJ_2(compare_not_equal_obj, compare_not_equal);
 
-//| def maximum(x1, x2):
+//| def maximum(x1: Union[ulab.array, float], x2: Union[ulab.array, float]) -> ulab.array:
-//|    """
+//|     """
-//|    Compute the element by element maximum of the arguments.
+//|     Compute the element by element maximum of the arguments.
-//|    Arguments may be ulab arrays or numbers.  All array arguments
+//|
-//|    must be the same size.  If the inputs are both scalars, a number is
+//|     Arguments may be ulab arrays or numbers.  All array arguments
-//|    returned"""
+//|     must be the same size.  If the inputs are both scalars, a number is
-//|    ...
+//|     returned"""
+//|     ...
 //|
-
 static mp_obj_t compare_maximum(mp_obj_t x1, mp_obj_t x2) {
 	// extra round, so that we can return maximum(3, 4) properly
 	mp_obj_t result = compare_function(x1, x2, COMPARE_MAXIMUM);
@@ -198,12 +201,13 @@ static mp_obj_t compare_maximum(mp_obj_t x1, mp_obj_t x2) {
 
 MP_DEFINE_CONST_FUN_OBJ_2(compare_maximum_obj, compare_maximum);
 
-//| def minimum(x1, x2):
+//| def minimum(x1: Union[ulab.array, float], x2: Union[ulab.array, float]) -> ulab.array:
-//|    """Compute the element by element minimum of the arguments.
+//|     """Compute the element by element minimum of the arguments.
-//|    Arguments may be ulab arrays or numbers.  All array arguments
+//|
-//|    must be the same size.  If the inputs are both scalars, a number is
+//|     Arguments may be ulab arrays or numbers.  All array arguments
-//|    returned"""
+//|     must be the same size.  If the inputs are both scalars, a number is
-//|    ...
+//|     returned"""
+//|     ...
 //|
 
 static mp_obj_t compare_minimum(mp_obj_t x1, mp_obj_t x2) {

code/fft/fft.c

@@ -156,13 +156,14 @@ mp_obj_t fft_fft_ifft_spectrum(size_t n_args, mp_obj_t arg_re, mp_obj_t arg_im,
     }
 }
 
-//| def fft(r, c=None):
+//| def fft(r: ulab.array, c: Optional[ulab.array] = None) -> Tuple[ulab.array, ulab.array]:
 //|     """
 //|     :param ulab.array r: A 1-dimension array of values whose size is a power of 2
 //|     :param ulab.array c: An optional 1-dimension array of values whose size is a power of 2, giving the complex part of the value
 //|     :return tuple (r, c): The real and complex parts of the FFT
 //|
 //|     Perform a Fast Fourier Transform from the time domain into the frequency domain
+//|
 //|     See also ~ulab.extras.spectrum, which computes the magnitude of the fft,
 //|     rather than separately returning its real and imaginary parts."""
 //|     ...
@@ -178,7 +179,7 @@ static mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
 
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_fft_obj, 1, 2, fft_fft);
 
-//| def ifft(r, c=None):
+//| def ifft(r: ulab.array, c: Optional[ulab.array] = None) -> Tuple[ulab.array, ulab.array]:
 //|     """
 //|     :param ulab.array r: A 1-dimension array of values whose size is a power of 2
 //|     :param ulab.array c: An optional 1-dimension array of values whose size is a power of 2, giving the complex part of the value
@@ -198,7 +199,7 @@ static mp_obj_t fft_ifft(size_t n_args, const mp_obj_t *args) {
 
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_ifft_obj, 1, 2, fft_ifft);
 
-//| def spectrogram(r):
+//| def spectrogram(r: ulab.array) -> ulab.array:
 //|     """
 //|     :param ulab.array r: A 1-dimension array of values whose size is a power of 2
 //|

code/linalg/linalg.c

@@ -80,14 +80,14 @@ bool linalg_invert_matrix(mp_float_t *data, size_t N) {
     return true;
 }
 
-//| def cholesky(A):
+//| def cholesky(A: ulab.array) -> ulab.array:
-//|    """
+//|     """
-//|    :param ~ulab.array A: a positive definite, symmetric square matrix
+//|     :param ~ulab.array A: a positive definite, symmetric square matrix
-//|    :return ~ulab.array L: a square root matrix in the lower triangular form
+//|     :return ~ulab.array L: a square root matrix in the lower triangular form
-//|    :raises ValueError: If the input does not fulfill the necessary conditions
+//|     :raises ValueError: If the input does not fulfill the necessary conditions
 //|
-//|    The returned matrix satisfies the equation m=LL*"""
+//|     The returned matrix satisfies the equation m=LL*"""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t linalg_cholesky(mp_obj_t oin) {
@@ -141,12 +141,13 @@ static mp_obj_t linalg_cholesky(mp_obj_t oin) {
 
 MP_DEFINE_CONST_FUN_OBJ_1(linalg_cholesky_obj, linalg_cholesky);
 
-//| def det():
+//| def det(m: ulab.array) -> float:
-//|    """
+//|     """
-//|    :param: m, a square matrix
+//|     :param: m, a square matrix
-//|    :return float: The determinant of the matrix"""
+//|     :return float: The determinant of the matrix
 //|
-//|    ...
+//|     Computes the eigenvalues and eigenvectors of a square matrix"""
+//|     ...
 //|
 
 static mp_obj_t linalg_det(mp_obj_t oin) {
@@ -181,7 +182,7 @@ static mp_obj_t linalg_det(mp_obj_t oin) {
 
 MP_DEFINE_CONST_FUN_OBJ_1(linalg_det_obj, linalg_det);
 
-//| def dot(m1, m2):
+//| def dot(m1: ulab.array, m2: ulab.array) -> Union[ulab.array, float]:
 //|    """
 //|    :param ~ulab.array m1: a matrix, or a vector
 //|    :param ~ulab.array m2: a matrix, or a vector
@@ -234,13 +235,13 @@ static mp_obj_t linalg_dot(mp_obj_t _m1, mp_obj_t _m2) {
 
 MP_DEFINE_CONST_FUN_OBJ_2(linalg_dot_obj, linalg_dot);
 
-//| def eig(m):
+//| def eig(m: ulab.array) -> Tuple[ulab.array, ulab.array]:
-//|    """
+//|     """
-//|    :param m: a square matrix
+//|     :param m: a square matrix
-//|    :return tuple (eigenvectors, eigenvalues):
+//|     :return tuple (eigenvectors, eigenvalues):
 //|
-//|    Computes the eigenvalues and eigenvectors of a square matrix"""
+//|     Computes the eigenvalues and eigenvectors of a square matrix"""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t linalg_eig(mp_obj_t oin) {
@@ -365,16 +366,15 @@ static mp_obj_t linalg_eig(mp_obj_t oin) {
 
 MP_DEFINE_CONST_FUN_OBJ_1(linalg_eig_obj, linalg_eig);
 
-//| def inv(m):
+//| def inv(m: ulab.array) -> ulab.array:
-//|    """
+//|     """
-//|    :param ~ulab.array m: a square matrix
+//|     :param ~ulab.array m: a square matrix
-//|    :return: The inverse of the matrix, if it exists
+//|     :return: The inverse of the matrix, if it exists
-//|    :raises ValueError: if the matrix is not invertible
+//|     :raises ValueError: if the matrix is not invertible
 //|
-//|    Computes the inverse of a square matrix"""
+//|     Computes the inverse of a square matrix"""
-//|    ...
+//|     ...
 //|
-
 static mp_obj_t linalg_inv(mp_obj_t o_in) {
 	ndarray_obj_t *o = linalg_object_is_square(o_in);
     ndarray_obj_t *inverted = create_new_ndarray(o->m, o->n, NDARRAY_FLOAT);
@@ -423,9 +423,9 @@ static mp_obj_t linalg_norm(mp_obj_t _x) {
 
 MP_DEFINE_CONST_FUN_OBJ_1(linalg_norm_obj, linalg_norm);
 
-//| def size(array):
+//| def size(array: ulab.array) -> int:
-//|    """Return the total number of elements in the array, as an integer."""
+//|     """Return the total number of elements in the array, as an integer."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t linalg_size(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -468,12 +468,12 @@ static mp_obj_t linalg_size(size_t n_args, const mp_obj_t *pos_args, mp_map_t *k
 
 MP_DEFINE_CONST_FUN_OBJ_KW(linalg_size_obj, 1, linalg_size);
 
-//| def trace(m):
+//| def trace(m: ulab.array) -> float:
-//|    """
+//|     """
-//|    :param m: a square matrix
+//|     :param m: a square matrix
 //|
-//|    Compute the trace of the matrix, the sum of its diagonal elements."""
+//|     Compute the trace of the matrix, the sum of its diagonal elements."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t linalg_trace(mp_obj_t oin) {

code/ndarray.c

@@ -36,9 +36,24 @@ mp_uint_t ndarray_print_edgeitems = NDARRAY_PRINT_EDGEITEMS;
 //| possible.  Numpy's documentation can be found at
 //| https://docs.scipy.org/doc/numpy/index.html"""
 //|
+//| from typing import List
+//|
+//| _DType = int
+//| """`ulab.int8`, `ulab.uint8`, `ulab.int16`, `ulab.uint16`, or `ulab.float`"""
+//|
+//| _Index = Union[int, slice, List[bool], Tuple[Union[int, slice, List[bool]], Union[int, slice, List[bool]]]]
+//| _float = float
+//|
+
 //| class array:
 //|     """1- and 2- dimensional array"""
-//|     def __init__(self, values, *, dtype=float):
+//|
+//|     def __init__(
+//|         self,
+//|         values: Union[array, Iterable[_float], Iterable[Iterable[_float]]],
+//|         *,
+//|         dtype: _DType = float
+//|     ) -> None:
 //|         """:param sequence values: Sequence giving the initial content of the array.
 //|           :param dtype: The type of array values, ``int8``, ``uint8``, ``int16``, ``uint16``, or ``float``
 //|
@@ -54,16 +69,16 @@ mp_uint_t ndarray_print_edgeitems = NDARRAY_PRINT_EDGEITEMS;
 //|           places an `array.array` can be used."""
 //|         ...
 //|
-//|     shape: tuple = ...
+//|     shape: Union[Tuple[int], Tuple[int, int]]
 //|     """The size of the array, a tuple of length 1 or 2"""
 //|
-//|     size: int = ...
+//|     size: int
 //|     """The number of elements in the array"""
 //|
-//|     itemsize: int = ...
+//|     itemsize: int
-//|     """The number of elements in the array"""
+//|     """The size of a single item in the array"""
 //|
-//|     def flatten(self, *, order='C'):
+//|     def flatten(self, *, order: str = "C") -> array:
 //|         """:param order: Whether to flatten by rows ('C') or columns ('F')
 //|
 //|            Returns a new `ulab.array` object which is always 1 dimensional.
@@ -72,63 +87,96 @@ mp_uint_t ndarray_print_edgeitems = NDARRAY_PRINT_EDGEITEMS;
 //|            to the typical storage organization of the C and Fortran languages."""
 //|         ...
 //|
-//|     def sort(self, *, axis=1):
+//|     def reshape(self, shape: Tuple[int, int]) -> array:
+//|         """Returns an array containing the same data with a new shape."""
+//|         ...
+//|
+//|     def sort(self, *, axis: Optional[int] = 1) -> None:
 //|         """:param axis: Whether to sort elements within rows (0), columns (1), or elements (None)"""
 //|         ...
 //|
-//|     def transpose(self):
+//|     def transpose(self) -> array:
 //|         """Swap the rows and columns of a 2-dimensional array"""
 //|         ...
 //|
-//|     def  __add__(self):
+//|     def __add__(self, other: Union[array, _float]) -> array:
 //|         """Adds corresponding elements of the two arrays, or adds a number to all
 //|            elements of the array.  If both arguments are arrays, their sizes must match."""
 //|         ...
+//|     def __radd__(self, other: _float) -> array: ...
 //|
-//|     def  __sub__(self):
+//|     def __sub__(self, other: Union[array, _float]) -> array:
-//|         """Subtracts corresponding elements of the two arrays, or adds a number to all
+//|         """Subtracts corresponding elements of the two arrays, or subtracts a number from all
 //|            elements of the array.  If both arguments are arrays, their sizes must match."""
 //|         ...
+//|     def __rsub__(self, other: _float) -> array: ...
 //|
-//|     def  __mul__(self):
+//|     def __mul__(self, other: Union[array, _float]) -> array:
 //|         """Multiplies corresponding elements of the two arrays, or multiplies
 //|            all elements of the array by a number.  If both arguments are arrays,
 //|            their sizes must match."""
 //|         ...
+//|     def __rmul__(self, other: _float) -> array: ...
 //|
-//|     def __div__(self):
+//|     def __div__(self, other: Union[array, _float]) -> array:
 //|         """Multiplies corresponding elements of the two arrays, or divides
 //|            all elements of the array by a number.  If both arguments are arrays,
 //|            their sizes must match."""
 //|         ...
+//|     def __rdiv__(self, other: _float) -> array: ...
 //|
-//|     def __pow__():
+//|     def __pow__(self, other: Union[array, _float]) -> array:
 //|         """Computes the power (x**y) of corresponding elements of the the two arrays,
 //|            or one number and one array.  If both arguments are arrays, their sizes
 //|            must match."""
 //|         ...
+//|     def __rpow__(self, other: _float) -> array: ...
 //|
-//|     def __getitem__():
+//|     def __inv__(self) -> array:
+//|         ...
+//|     def __neg__(self) -> array:
+//|         ...
+//|     def __pos__(self) -> array:
+//|         ...
+//|     def __abs__(self) -> array:
+//|         ...
+//|     def __len__(self) -> array:
+//|         ...
+//|     def __lt__(self, other: Union[array, _float]) -> List[bool]:
+//|         ...
+//|     def __le__(self, other: Union[array, _float]) -> List[bool]:
+//|         ...
+//|     def __gt__(self, other: Union[array, _float]) -> List[bool]:
+//|         ...
+//|     def __ge__(self, other: Union[array, _float]) -> List[bool]:
+//|         ...
+//|
+//|     def __iter__(self) -> Union[Iterator[array], Iterator[_float]]:
+//|         ...
+//|
+//|     def __getitem__(self, index: _Index) -> Union[array, _float]:
 //|         """Retrieve an element of the array."""
 //|         ...
 //|
-//|     def __setitem__():
+//|     def __setitem__(self, index: _Index, value: Union[array, _float]) -> None:
 //|         """Set an element of the array."""
 //|         ...
 //|
-//| int8 = ...
+//| _ArrayLike = Union[array, List[_float], Tuple[_float], range]
+//|
+//| int8: _DType
 //| """Type code for signed integers in the range -128 .. 127 inclusive, like the 'b' typecode of `array.array`"""
 //|
-//| int16 = ...
+//| int16: _DType
 //| """Type code for signed integers in the range -32768 .. 32767 inclusive, like the 'h' typecode of `array.array`"""
 //|
-//| float = ...
+//| float: _DType
 //| """Type code for floating point values, like the 'f' typecode of `array.array`"""
 //|
-//| uint8 = ...
+//| uint8: _DType
 //| """Type code for unsigned integers in the range 0 .. 255 inclusive, like the 'H' typecode of `array.array`"""
 //|
-//| uint16 = ...
+//| uint16: _DType
 //| """Type code for unsigned integers in the range 0 .. 65535 inclusive, like the 'h' typecode of `array.array`"""
 //|
 

code/numerical/numerical.c

@@ -36,6 +36,9 @@ enum NUMERICAL_FUNCTION_TYPE {
 //|
 //| Most of these functions take an "axis" argument, which indicates whether to
 //| operate over the flattened array (None), rows (0), or columns (1)."""
+//|
+//| from ulab import _ArrayLike
+//|
 
 static void axis_sorter(ndarray_obj_t *ndarray, mp_obj_t axis, size_t *m, size_t *n, size_t *N, 
                  size_t *increment, size_t *len, size_t *start_inc) {
@@ -340,9 +343,9 @@ static mp_obj_t numerical_sort_helper(mp_obj_t oin, mp_obj_t axis, uint8_t inpla
     }
 }
 
-//| def argmax(array, *, axis=None):
+//| def argmax(array: _ArrayLike, *, axis: Optional[int] = None) -> int:
-//|    """Return the index of the maximum element of the 1D array"""
+//|     """Return the index of the maximum element of the 1D array"""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_argmax(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -351,9 +354,9 @@ static mp_obj_t numerical_argmax(size_t n_args, const mp_obj_t *pos_args, mp_map
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_argmax_obj, 1, numerical_argmax);
 
-//| def argmin(array, *, axis=None):
+//| def argmin(array: _ArrayLike, *, axis: Optional[int] = None) -> int:
-//|    """Return the index of the minimum element of the 1D array"""
+//|     """Return the index of the minimum element of the 1D array"""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_argmin(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -362,9 +365,9 @@ static mp_obj_t numerical_argmin(size_t n_args, const mp_obj_t *pos_args, mp_map
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_argmin_obj, 1, numerical_argmin);
 
-//| def argsort(array, *, axis=None):
+//| def argsort(array: ulab.array, *, axis: Optional[int] = None) -> ulab.array:
-//|    """Returns an array which gives indices into the input array from least to greatest."""
+//|     """Returns an array which gives indices into the input array from least to greatest."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_argsort(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -437,10 +440,10 @@ static mp_obj_t numerical_argsort(size_t n_args, const mp_obj_t *pos_args, mp_ma
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_argsort_obj, 1, numerical_argsort);
 
-//| def diff(array, *, axis=1):
+//| def diff(array: ulab.array, *, axis: int = 1) -> ulab.array:
-//|    """Return the numerical derivative of successive elements of the array, as
+//|     """Return the numerical derivative of successive elements of the array, as
-//|       an array.  axis=None is not supported."""
+//|        an array.  axis=None is not supported."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_diff(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -513,10 +516,10 @@ static mp_obj_t numerical_diff(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_diff_obj, 1, numerical_diff);
 
-//| def flip(array, *, axis=None):
+//| def flip(array: ulab.array, *, axis: Optional[int] = None) -> ulab.array:
-//|    """Returns a new array that reverses the order of the elements along the
+//|     """Returns a new array that reverses the order of the elements along the
-//|       given axis, or along all axes if axis is None."""
+//|        given axis, or along all axes if axis is None."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_flip(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -568,9 +571,9 @@ static mp_obj_t numerical_flip(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_flip_obj, 1, numerical_flip);
 
-//| def max(array, *, axis=None):
+//| def max(array: _ArrayLike, *, axis: Optional[int] = None) -> float:
-//|    """Return the maximum element of the 1D array"""
+//|     """Return the maximum element of the 1D array"""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_max(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -579,9 +582,9 @@ static mp_obj_t numerical_max(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_max_obj, 1, numerical_max);
 
-//| def mean(array, *, axis=None):
+//| def mean(array: _ArrayLike, *, axis: Optional[int] = None) -> float:
-//|    """Return the mean element of the 1D array, as a number if axis is None, otherwise as an array."""
+//|     """Return the mean element of the 1D array, as a number if axis is None, otherwise as an array."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_mean(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -590,9 +593,9 @@ static mp_obj_t numerical_mean(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_mean_obj, 1, numerical_mean);
 
-//| def min(array, *, axis=None):
+//| def min(array: _ArrayLike, *, axis: Optional[int] = None) -> float:
-//|    """Return the minimum element of the 1D array"""
+//|     """Return the minimum element of the 1D array"""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_min(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -601,11 +604,11 @@ static mp_obj_t numerical_min(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_min_obj, 1, numerical_min);
 
-//| def roll(array, distance, *, axis=None):
+//| def roll(array: ulab.array, distance: int, *, axis: Optional[int] = None) -> None:
-//|    """Shift the content of a vector by the positions given as the second
+//|     """Shift the content of a vector by the positions given as the second
-//|       argument. If the ``axis`` keyword is supplied, the shift is applied to
+//|        argument. If the ``axis`` keyword is supplied, the shift is applied to
-//|       the given axis.  The array is modified in place."""
+//|        the given axis.  The array is modified in place."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_roll(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -690,10 +693,10 @@ static mp_obj_t numerical_roll(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_roll_obj, 2, numerical_roll);
 
-//| def sort(array, *, axis=0):
+//| def sort(array: ulab.array, *, axis: Optional[int] = 0) -> ulab.array:
-//|    """Sort the array along the given axis, or along all axes if axis is None.
+//|     """Sort the array along the given axis, or along all axes if axis is None.
-//|       The array is modified in place."""
+//|        The array is modified in place."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_sort(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -725,9 +728,9 @@ static mp_obj_t numerical_sort_inplace(size_t n_args, const mp_obj_t *pos_args,
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_sort_inplace_obj, 1, numerical_sort_inplace);
 
-//| def std(array, *, axis=None):
+//| def std(array: _ArrayLike, *, axis: Optional[int] = None) -> float:
-//|    """Return the standard deviation of the array, as a number if axis is None, otherwise as an array."""
+//|     """Return the standard deviation of the array, as a number if axis is None, otherwise as an array."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_std(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -760,9 +763,9 @@ static mp_obj_t numerical_std(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
 MP_DEFINE_CONST_FUN_OBJ_KW(numerical_std_obj, 1, numerical_std);
 
-//| def sum(array, *, axis=None):
+//| def sum(array: _ArrayLike, *, axis: Optional[int] = None) -> Union[float, int, ulab.array]:
-//|    """Return the sum of the array, as a number if axis is None, otherwise as an array."""
+//|     """Return the sum of the array, as a number if axis is None, otherwise as an array."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t numerical_sum(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {

code/poly/poly.c

@@ -21,11 +21,16 @@
 
 //| """Polynomial functions"""
 //|
+//| from ulab import _ArrayLike
+//|
 
-//| def polyfit(x, y, degree):
+//| @overload
-//|    """Return a polynomial of given degree that approximates the function
+//| def polyfit(y: _ArrayLike, degree: int) -> ulab.array: ...
-//|       f(x)=y.  If x is not supplied, it is the range(len(y))."""
+//| @overload
-//|    ...
+//| def polyfit(x: _ArrayLike, y: _ArrayLike, degree: int) -> ulab.array:
+//|     """Return a polynomial of given degree that approximates the function
+//|        f(x)=y.  If x is not supplied, it is the range(len(y))."""
+//|     ...
 //|
 
 static mp_obj_t poly_polyfit(size_t  n_args, const mp_obj_t *args) {
@@ -143,9 +148,9 @@ static mp_obj_t poly_polyfit(size_t  n_args, const mp_obj_t *args) {
 
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(poly_polyfit_obj, 2, 3, poly_polyfit);
 
-//| def polyval(p, x):
+//| def polyval(p: _ArrayLike, x: _ArrayLike) -> ulab.array:
-//|    """Evaluate the polynomial p at the points x.  x must be an array."""
+//|     """Evaluate the polynomial p at the points x.  x must be an array."""
-//|    ...
+//|     ...
 //|
 
 static mp_obj_t poly_polyval(mp_obj_t o_p, mp_obj_t o_x) {

code/ulab_create.c

@@ -69,19 +69,22 @@ STATIC ndarray_obj_t *create_linspace_arange(mp_float_t start, mp_float_t step,
     return ndarray;
 }
 
-//| def arange(start, stop, step, dtype=float):
+//| @overload
-//|    """
+//| def arange(stop: _float, step: _float = 1, dtype: _DType = float) -> array: ...
-//|    .. param: start
+//| @overload
-//|      First value in the array, optional, defaults to 0
+//| def arange(start: _float, stop: _float, step: _float = 1, dtype: _DType = float) -> array:
-//|    .. param: stop
+//|     """
-//|      Final value in the array
+//|     .. param: start
-//|    .. param: step
+//|       First value in the array, optional, defaults to 0
-//|      Difference between consecutive elements, optional, defaults to 1.0
+//|     .. param: stop
-//|    .. param: dtype
+//|       Final value in the array
-//|      Type of values in the array
+//|     .. param: step
+//|       Difference between consecutive elements, optional, defaults to 1.0
+//|     .. param: dtype
+//|       Type of values in the array
 //|
-//|    Return a new 1-D array with elements ranging from ``start`` to ``stop``, with step size ``step``."""
+//|     Return a new 1-D array with elements ranging from ``start`` to ``stop``, with step size ``step``."""
-//|    ...
+//|     ...
 //|
 
 mp_obj_t create_arange(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -131,12 +134,11 @@ mp_obj_t create_arange(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_arg
 
 MP_DEFINE_CONST_FUN_OBJ_KW(create_arange_obj, 1, create_arange);
 
-//| def eye(size, *, dtype=float):
+//| def eye(size: int, *, dtype: _DType = float) -> array:
-//|    """Return a new square array of size, with the diagonal elements set to 1
+//|     """Return a new square array of size, with the diagonal elements set to 1
-//|       and the other elements set to 0."""
+//|        and the other elements set to 0."""
-//|    ...
+//|     ...
 //|
-
 /*
 mp_obj_t create_eye(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     static const mp_arg_t allowed_args[] = {
@@ -180,25 +182,31 @@ mp_obj_t create_eye(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args)
 }
 
 MP_DEFINE_CONST_FUN_OBJ_KW(create_eye_obj, 0, create_eye);
-
 */
-//| def linspace(start, stop, *, dtype=float, num=50, endpoint=True):
+//| def linspace(
-//|    """
+//|     start: _float,
-//|    .. param: start
+//|     stop: _float,
-//|      First value in the array
+//|     *,
-//|    .. param: stop
+//|     dtype: _DType = float,
-//|      Final value in the array
+//|     num: int = 50,
-//|    .. param int: num
+//|     endpoint: bool = True
-//|      Count of values in the array
+//| ) -> array:
-//|    .. param: dtype
+//|     """
-//|      Type of values in the array
+//|     .. param: start
-//|    .. param bool: endpoint
+//|       First value in the array
-//|      Whether the ``stop`` value is included.  Note that even when
+//|     .. param: stop
-//|      endpoint=True, the exact ``stop`` value may not be included due to the
+//|       Final value in the array
-//|      inaccuracy of floating point arithmetic.
+//|     .. param int: num
+//|       Count of values in the array
+//|     .. param: dtype
+//|       Type of values in the array
+//|     .. param bool: endpoint
+//|       Whether the ``stop`` value is included.  Note that even when
+//|       endpoint=True, the exact ``stop`` value may not be included due to the
+//|       inaccuracy of floating point arithmetic.
 //|
-//|    Return a new 1-D array with ``num`` elements ranging from ``start`` to ``stop`` linearly."""
+//|     Return a new 1-D array with ``num`` elements ranging from ``start`` to ``stop`` linearly."""
-//|    ...
+//|     ...
 //|
 
 mp_obj_t create_linspace(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -236,7 +244,7 @@ mp_obj_t create_linspace(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_a
 
 MP_DEFINE_CONST_FUN_OBJ_KW(create_linspace_obj, 2, create_linspace);
 
-//| def ones(shape, *, dtype=float):
+//| def ones(shape: Union[int, Tuple[int, int]], *, dtype: _DType = float) -> array:
 //|    """
 //|    .. param: shape
 //|       Shape of the array, either an integer (for a 1-D array) or a tuple of 2 integers (for a 2-D array)
@@ -253,7 +261,7 @@ mp_obj_t create_ones(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args)
 
 MP_DEFINE_CONST_FUN_OBJ_KW(create_ones_obj, 0, create_ones);
 
-//| def zeros(shape, *, dtype):
+//| def zeros(shape: Union[int, Tuple[int, int]], *, dtype: _DType = float) -> array:
 //|    """
 //|    .. param: shape
 //|       Shape of the array, either an integer (for a 1-D array) or a tuple of 2 integers (for a 2-D array)

code/vector/vectorise.c

@@ -31,6 +31,8 @@
 //| applying the function to every element in the array.  This is typically
 //| much more efficient than expressing the same operation as a Python loop."""
 //|
+//| from ulab import _DType, _ArrayLike
+//|
     
 static mp_obj_t vectorise_generic_vector(mp_obj_t o_in, mp_float_t (*f)(mp_float_t)) {
     // Return a single value, if o_in is not iterable
@@ -71,7 +73,7 @@ static mp_obj_t vectorise_generic_vector(mp_obj_t o_in, mp_float_t (*f)(mp_float
     return mp_const_none;
 }
 
-//| def acos():
+//| def acos(a: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse cosine function"""
 //|    ...
 //|
@@ -79,7 +81,7 @@ static mp_obj_t vectorise_generic_vector(mp_obj_t o_in, mp_float_t (*f)(mp_float
 MATH_FUN_1(acos, acos);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_acos_obj, vectorise_acos);
 
-//| def acosh():
+//| def acosh(a: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse hyperbolic cosine function"""
 //|    ...
 //|
@@ -88,7 +90,7 @@ MATH_FUN_1(acosh, acosh);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_acosh_obj, vectorise_acosh);
 
 
-//| def asin():
+//| def asin(a: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse sine function"""
 //|    ...
 //|
@@ -96,7 +98,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_acosh_obj, vectorise_acosh);
 MATH_FUN_1(asin, asin);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_asin_obj, vectorise_asin);
 
-//| def asinh():
+//| def asinh(a: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse hyperbolic sine function"""
 //|    ...
 //|
@@ -105,7 +107,7 @@ MATH_FUN_1(asinh, asinh);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_asinh_obj, vectorise_asinh);
 
 
-//| def around(a, *, decimals):
+//| def around(a: _ArrayLike, *, decimals: int = 0) -> ulab.array:
 //|    """Returns a new float array in which each element is rounded to
 //|       ``decimals`` places."""
 //|    ...
@@ -136,7 +138,7 @@ static mp_obj_t vectorise_around(size_t n_args, const mp_obj_t *pos_args, mp_map
 
 MP_DEFINE_CONST_FUN_OBJ_KW(vectorise_around_obj, 1, vectorise_around);
 
-//| def atan():
+//| def atan(a: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse tangent function; the return values are in the
 //|       range [-pi/2,pi/2]."""
 //|    ...
@@ -145,7 +147,7 @@ MP_DEFINE_CONST_FUN_OBJ_KW(vectorise_around_obj, 1, vectorise_around);
 MATH_FUN_1(atan, atan);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_atan_obj, vectorise_atan);
 
-//| def atan2(y,x):
+//| def atan2(ya: _ArrayLike, xa: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse tangent function of y/x; the return values are in
 //|       the range [-pi, pi]."""
 //|    ...
@@ -206,7 +208,7 @@ MP_DEFINE_CONST_FUN_OBJ_2(vectorise_arctan2_obj, vectorise_arctan2);
 
 
 
-//| def atanh():
+//| def atanh(a: _ArrayLike) -> ulab.array:
 //|    """Computes the inverse hyperbolic tangent function"""
 //|    ...
 //|
@@ -214,7 +216,7 @@ MP_DEFINE_CONST_FUN_OBJ_2(vectorise_arctan2_obj, vectorise_arctan2);
 MATH_FUN_1(atanh, atanh);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_atanh_obj, vectorise_atanh);
 
-//| def ceil():
+//| def ceil(a: _ArrayLike) -> ulab.array:
 //|    """Rounds numbers up to the next whole number"""
 //|    ...
 //|
@@ -222,7 +224,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_atanh_obj, vectorise_atanh);
 MATH_FUN_1(ceil, ceil);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_ceil_obj, vectorise_ceil);
 
-//| def cos():
+//| def cos(a: _ArrayLike) -> ulab.array:
 //|    """Computes the cosine function"""
 //|    ...
 //|
@@ -230,7 +232,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_ceil_obj, vectorise_ceil);
 MATH_FUN_1(cos, cos);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_cos_obj, vectorise_cos);
 
-//| def cosh():
+//| def cosh(a: _ArrayLike) -> ulab.array:
 //|    """Computes the hyperbolic cosine function"""
 //|    ...
 //|
@@ -238,7 +240,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_cos_obj, vectorise_cos);
 MATH_FUN_1(cosh, cosh);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_cosh_obj, vectorise_cosh);
 
-//| def erf():
+//| def erf(a: _ArrayLike) -> ulab.array:
 //|    """Computes the error function, which has applications in statistics"""
 //|    ...
 //|
@@ -246,7 +248,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_cosh_obj, vectorise_cosh);
 MATH_FUN_1(erf, erf);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_erf_obj, vectorise_erf);
 
-//| def erfc():
+//| def erfc(a: _ArrayLike) -> ulab.array:
 //|    """Computes the complementary error function, which has applications in statistics"""
 //|    ...
 //|
@@ -254,7 +256,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_erf_obj, vectorise_erf);
 MATH_FUN_1(erfc, erfc);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_erfc_obj, vectorise_erfc);
 
-//| def exp():
+//| def exp(a: _ArrayLike) -> ulab.array:
 //|    """Computes the exponent function."""
 //|    ...
 //|
@@ -262,7 +264,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_erfc_obj, vectorise_erfc);
 MATH_FUN_1(exp, exp);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_exp_obj, vectorise_exp);
 
-//| def expm1():
+//| def expm1(a: _ArrayLike) -> ulab.array:
 //|    """Computes $e^x-1$.  In certain applications, using this function preserves numeric accuracy better than the `exp` function."""
 //|    ...
 //|
@@ -270,7 +272,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_exp_obj, vectorise_exp);
 MATH_FUN_1(expm1, expm1);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_expm1_obj, vectorise_expm1);
 
-//| def floor():
+//| def floor(a: _ArrayLike) -> ulab.array:
 //|    """Rounds numbers up to the next whole number"""
 //|    ...
 //|
@@ -278,7 +280,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_expm1_obj, vectorise_expm1);
 MATH_FUN_1(floor, floor);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_floor_obj, vectorise_floor);
 
-//| def gamma():
+//| def gamma(a: _ArrayLike) -> ulab.array:
 //|    """Computes the gamma function"""
 //|    ...
 //|
@@ -286,7 +288,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_floor_obj, vectorise_floor);
 MATH_FUN_1(gamma, tgamma);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_gamma_obj, vectorise_gamma);
 
-//| def lgamma():
+//| def lgamma(a: _ArrayLike) -> ulab.array:
 //|    """Computes the natural log of the gamma function"""
 //|    ...
 //|
@@ -294,7 +296,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_gamma_obj, vectorise_gamma);
 MATH_FUN_1(lgamma, lgamma);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_lgamma_obj, vectorise_lgamma);
 
-//| def log():
+//| def log(a: _ArrayLike) -> ulab.array:
 //|    """Computes the natural log"""
 //|    ...
 //|
@@ -302,7 +304,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_lgamma_obj, vectorise_lgamma);
 MATH_FUN_1(log, log);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_log_obj, vectorise_log);
 
-//| def log10():
+//| def log10(a: _ArrayLike) -> ulab.array:
 //|    """Computes the log base 10"""
 //|    ...
 //|
@@ -310,7 +312,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_log_obj, vectorise_log);
 MATH_FUN_1(log10, log10);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_log10_obj, vectorise_log10);
 
-//| def log2():
+//| def log2(a: _ArrayLike) -> ulab.array:
 //|    """Computes the log base 2"""
 //|    ...
 //|
@@ -318,7 +320,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_log10_obj, vectorise_log10);
 MATH_FUN_1(log2, log2);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_log2_obj, vectorise_log2);
 
-//| def sin():
+//| def sin(a: _ArrayLike) -> ulab.array:
 //|    """Computes the sine function"""
 //|    ...
 //|
@@ -326,7 +328,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_log2_obj, vectorise_log2);
 MATH_FUN_1(sin, sin);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_sin_obj, vectorise_sin);
 
-//| def sinh():
+//| def sinh(a: _ArrayLike) -> ulab.array:
 //|    """Computes the hyperbolic sine"""
 //|    ...
 //|
@@ -334,7 +336,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_sin_obj, vectorise_sin);
 MATH_FUN_1(sinh, sinh);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_sinh_obj, vectorise_sinh);
 
-//| def sqrt():
+//| def sqrt(a: _ArrayLike) -> ulab.array:
 //|    """Computes the square root"""
 //|    ...
 //|
@@ -342,7 +344,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_sinh_obj, vectorise_sinh);
 MATH_FUN_1(sqrt, sqrt);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_sqrt_obj, vectorise_sqrt);
 
-//| def tan():
+//| def tan(a: _ArrayLike) -> ulab.array:
 //|    """Computes the tangent"""
 //|    ...
 //|
@@ -350,7 +352,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(vectorise_sqrt_obj, vectorise_sqrt);
 MATH_FUN_1(tan, tan);
 MP_DEFINE_CONST_FUN_OBJ_1(vectorise_tan_obj, vectorise_tan);
 
-//| def tanh():
+//| def tanh(a: _ArrayLike) -> ulab.array:
 //|    """Computes the hyperbolic tangent"""
 //|    ...
 
@@ -402,7 +404,11 @@ const mp_obj_type_t vectorise_function_type = {
     .call = vectorise_vectorized_function_call,
 };
 
-//| def vectorize(f, *, otypes=None):
+//| def vectorize(
+//|     f: Union[Callable[[int], float], Callable[[float], float]],
+//|     *,
+//|     otypes: Optional[_DType] = None
+//| ) -> Callable[[_ArrayLike], ulab.array]:
 //|    """
 //|    :param callable f: The function to wrap
 //|    :param otypes: List of array types that may be returned by the function.  None is interpreted to mean the return value is float.