# Flexible specialization on array shapes

**URL:** <https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936>\
**Category:** Support: How do I do ...?\
**Created:** [September 24, 2021, 4:10pm UTC](https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936 "2021-09-24T16:10:16Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![tmfrln](https://avatars.discourse-cdn.com/v4/letter/t/c68b51/32.png) [@tmfrln](https://numba.discourse.group/u/tmfrln)\
**Post date:** [September 24, 2021, 4:10pm UTC](https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936/1 "2021-09-24T16:10:16Z")

</div>

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:  
_c_ij = _a_i _b_j

- if _a_ and _b_ are matrices, the result should be a 4th order array:  
_c_ijkl = _a_ij _b_kl

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

```auto
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 🙂

Kind regards  
Tim

---

<div class="post-metadata">

**Author:** ![luk-f-a](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/luk-f-a/32/56_2.png) [@luk-f-a](https://numba.discourse.group/u/luk-f-a)\
**Post date:** [September 25, 2021, 2:25pm UTC](https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936/2 "2021-09-25T14:25:19Z")

</div>

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.

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

---

<div class="post-metadata">

**Author:** ![tmfrln](https://avatars.discourse-cdn.com/v4/letter/t/c68b51/32.png) [@tmfrln](https://numba.discourse.group/u/tmfrln)\
**Post date:** [September 26, 2021, 11:23am UTC](https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936/3 "2021-09-26T11:23:35Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![luk-f-a](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/luk-f-a/32/56_2.png) [@luk-f-a](https://numba.discourse.group/u/luk-f-a)\
**Post date:** [September 26, 2021, 3:19pm UTC](https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936/4 "2021-09-26T15:19:31Z")

</div>

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:

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

```

---

<div class="post-metadata">

**Author:** ![tmfrln](https://avatars.discourse-cdn.com/v4/letter/t/c68b51/32.png) [@tmfrln](https://numba.discourse.group/u/tmfrln)\
**Post date:** [September 27, 2021, 7:13am UTC](https://numba.discourse.group/t/flexible-specialization-on-array-shapes/936/5 "2021-09-27T07:13:30Z")

</div>

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