# Flexible specialization on array shapes

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
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
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

@nb.generated_jit(nopython=True)
if (    isinstance(a, nb.types.Array) and a.ndim == 1
and isinstance(b, nb.types.Array) and b.ndim == 1):
# 4th order tensors

elif (    isinstance(a, nb.types.Array) and a.ndim == 2
and isinstance(b, nb.types.Array) and b.ndim == 2):
# 2nd order tensors / matrices

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)

``````

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

Kind regards
Tim

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

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

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.

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):