Flexible specialization on array shapes

Numba Support: How do I do ...?
1

Hi folks,

i recently started experimenting with numba and have a question considering the flexible specializations obtained via the @generated_jit decorator.

As an example, i want to generate a jit’ed function for the dyadic product of two array a and b:

  • if a and b are vectors, the result should be a matrix:
    cij = ai bj

  • if a and b are matrices, the result should be a 4th order array:
    cijkl = aij bkl

A working solution that i came up with is the following:

import numpy as np
import numba as nb

@nb.njit
def _t2_dyad_t2(a,b):
    dim = a.shape[0]

    c = np.zeros((dim,dim,dim,dim))
    for i in range(dim):
        for j in range(dim):
            for k in range(dim):
                for l in range(dim):
                    c[i,j,k,l] = a[i,j]*b[k,l]

    return c

@nb.njit
def _t1_dyad_t1(a,b):
    dim = a.shape[0]
    
    c = np.zeros((dim,dim))
    for i in range(dim):
        for j in range(dim):
            c[i,j] = a[i]*b[j]

    return c


# generate auto-specializations of dyad()
@nb.generated_jit(nopython=True)
def dyad(a,b):
    if (    isinstance(a, nb.types.Array) and a.ndim == 1 
        and isinstance(b, nb.types.Array) and b.ndim == 1):
        # 4th order tensors        
        return _t1_dyad_t1
    
    elif (    isinstance(a, nb.types.Array) and a.ndim == 2  
          and isinstance(b, nb.types.Array) and b.ndim == 2):
        # 2nd order tensors / matrices        
        return _t2_dyad_t2
    
    else:
        raise Exception("Dyadic product not implemented for this types.")
    

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

A = np.random.rand(3,3)
B = np.random.rand(3,3)

c = dyad(a,b)
C = dyad(A,B)

My question is if there is a shorter way to specify "a is a one-dimensional array", something like nb.types.Array[1] ?

On a side note, if there is an elegant way to generalize this type of product to arbitrary array combinations, that would be interesting too :slight_smile:

Kind regards
Tim

2

hi @tmfrln , when working with generated_jit the main body of the function is code that operates on types. You can use the property ndim of the array type to distinguish between your cases.

import numba
import numpy as np

@numba.generated_jit
def foo(x):
    if x.ndim == 1:
        return lambda x:"I'm a 1d array"
    elif x.ndim == 2:
        return lambda x:"I'm a 2d array"

foo(np.ones(1))

hope this helps
Luk

3

Hi @luk-f-a, thanks for the input :slight_smile:

While i basically did the same thing in my code above, your example made me realize that i do not need to actually query if the input is an array when i do not care about the case where it is not one.

4

hi @tmfrln, sorry I hadn’t realized that your code scrolled down, and I had only seen the first two functions, without the generated_jit part. I understand your question better now.

The only way I know how to write “an array with n dimensions” also includes specifying the type of the content and the order (C or F). numba.types.float64[:] is 1-d float array array(float64, 1d, A) and numba.types.int64[::1] is array(int64, 1d, C).

Checking isinstance(a, nb.types.Array) is not strictly necessary but it can be useful if you want to return a meaningful error message in case an unsupported type is passed. For example:

if not isinstance(a, nb.types.Array):
    raise TypeError("dyad only supports arrays")
5

Don’t worry - your answer was helpful anysways :wink:
Thanks again for the newbie support!