One of the main limitations of universal functions is that they must be defined on scalar values. A generalized universal function, abbreviated gufunc, is an extension of universal functions to procedures that take arrays.

A classic example is the matrix multiplication. In NumPy, matrix multiplication can be applied using the np.matmul function, which takes two 2D arrays and returns another 2D array. An example usage of np.matmul is as follows:

    a = np.random.rand(3, 3)
    b = np.random.rand(3, 3)

    c = np.matmul(a, b)
    c.shape
    # Result:
    # (3, 3)

As we saw in the previous subsection, a ufunc broadcasts the operation over arrays of scalars, its natural generalization will be to broadcast over an array of arrays. If, for instance, we take two arrays of 3 by 3 matrices, we will expect np.matmul to take to match the matrices and take their product. In the following example, we take two arrays containing 10 matrices of shape (3, 3). If we apply np.matmul, the product will be applied matrix-wise to obtain a new array containing the 10 results (which are, again, (3, 3) matrices):

    a = np.random.rand(10, 3, 3)
    b = np.random.rand(10, 3, 3)

    c = np.matmul(a, b)
    c.shape
    # Output
    # (10, 3, 3)

The usual rules for broadcasting will work in a similar way. For example, if we have an array of (3, 3) matrices, which will have a shape of (10, 3, 3), we can use np.matmul to calculate the matrix multiplication of each element with a single (3, 3) matrix. According to the broadcasting rules, we obtain that the single matrix will be repeated to obtain a size of (10, 3, 3):

    a = np.random.rand(10, 3, 3)
    b = np.random.rand(3, 3) # Broadcasted to shape (10, 3, 3)
    c = np.matmul(a, b)
    c.shape
    # Result:
    # (10, 3, 3)

Numba supports the implementation of efficient generalized universal functions through the nb.guvectorize decorator. As an example, we will implement a function that computes the euclidean distance between two arrays as a gufunc. To create a gufunc, we have to define a function that takes the input arrays, plus an output array where we will store the result of our calculation.

The nb.guvectorize decorator requires two arguments:

• The types of the input and output: two 1D arrays as input and a scalar as output
• The so called layout string, which is a representation of the input and output sizes; in our case, we take two arrays of the same size (denoted arbitrarily by n), and we output a scalar

In the following example, we show the implementation of the euclidean function using the nb.guvectorize decorator:

    @nb.guvectorize(['float64[:], float64[:], float64[:]'], '(n), (n) -
    > ()')
    def euclidean(a, b, out):
        N = a.shape[0]
        out[0] = 0.0
        for i in range(N):
            out[0] += (a[i] - b[i])**2

There are a few very important points to be made. Predictably, we declared the types of the inputs a and b as float64[:], because they are 1D arrays. However, what about the output argument? Wasn't it supposed to be a scalar? Yes, however, Numba treats scalar argument as arrays of size 1. That's why it was declared as float64[:].

Similarly, the layout string indicates that we have two arrays of size (n) and the output is a scalar, denoted by empty brackets--(). However, the array out will be passed as an array of size 1.

Also, note that we don't return anything from the function; all the output has to be written in the out array.

Our brand new euclidean function can be conveniently used on arrays of different shapes, as shown in the following example:

    a = np.random.rand(2)
    b = np.random.rand(2)
    c = euclidean(a, b) # Shape: (1,)

    a = np.random.rand(10, 2)
    b = np.random.rand(10, 2)
    c = euclidean(a, b) # Shape: (10,)

    a = np.random.rand(10, 2)
    b = np.random.rand(2)
    c = euclidean(a, b) # Shape: (10,)

How does the speed of euclidean compare to standard NumPy? In the following code, we benchmark a NumPy vectorized version with our previously defined euclidean function:

    a = np.random.rand(10000, 2)
    b = np.random.rand(10000, 2)

    %timeit ((a - b)**2).sum(axis=1)
    1000 loops, best of 3: 288 µs per loop

    %timeit euclidean(a, b)
    10000 loops, best of 3: 35.6 µs per loop

The Numba version, again, beats the NumPy version by a large margin!