List of different sized arrays as a parameter

Numba Community Support
1

I created a list of numpy arrays of different sizes: k = [ 1, np.ones(2), np.ones((2,2)), np.ones((2,2,2)), … ]. @njit understandably doesn’t like this but I’d like to be able to pass this as an argument. Is there pattern for this? One messy approach might be to flatten and concatenate these and rewrite everything in terms of slices or resize them after passing the 1D array.

@njit
def f(k):
    pass

@njit
def g(n):
    k = [ np.ones([2]*i) for i in range(n) ]
    f(k)
g(6)

TypeError: can't unbox heterogeneous list: array(float64, 1d, C) != array(float64, 2d, C)

2

This is hitting several limitations of Numba:

  • Numba treats arrays with difference number of dimensions as different types.
  • The list only support homogeneous types; i.e. [1d_array, 2d_array] would be heterogeneous.
  • Variables need to be type-stable; i.e for x in iterable cannot have x being multiple types (can’t be 1d_array and 2d_array in different iteration).

When you have dynamic types like that, the only solution I can think of is to not use @njit on g(). Instead write:

@njit
def f(k):
    ....


def g(n):
   for i in range(n):
       k = np.ones([2]*i)
       f(k)

g(6)

That way, f() is specialized on the different argument types.

1 Like
3

Thanks for elaborating on what was causing the issue. The problem is that I need to pass this list of k values since it’s the kernel for a multidimensional convolution. It’s a volterra series operator. The basis of that is that I have something like this k[0] + conv_1d1o(k[1], x) + conv_1d2o(k[2], x, x) + … Maybe I could generalize these convolutions better. With your suggested approach I’d need to continually recalculate k every time I were to use it. Additionally, I’d like to have @njit applied to g(n). Another option is to pass each one separately as an argument f(k[0], k[1], …) for every function that I pass k to and just limit the generalization to maybe 3 or so k’s. Each function in f has @njit applied, so maybe that’s the best I could have hoped for.

4

if you can create k as a tuple of arrays in interpreted code, and pass it to g() already as a tuple you can write

def g(ks):
   for k in literal_unroll(ks):
       f(k)
1 Like
5

Not sure if this would work for your case, but maybe the arrays can be created to have the maximum number of dimension by padding the leading dimension with ones. For example, if max ndim is 3, 1d array would have shape of (1, 1, N) and 2d array would have shape of (1, M, N).

1 Like
6

@sklam’s solution above is probably the best.

1 Like