How to set numba signature with nested lists?

Question

I'm trying to return an nested list, however running into some conversion error. Below is small piece of code for reproduction of error.

from numba import njit, prange

@njit("ListType(ListType(ListType(int32)))(int32, int32)", fastmath = True, parallel = True, cache = True)
def test(x, y):
    a = []
    for i in prange(10):
        b = []
        for j in range(4):
            c = []
            for k in range(5):
                c.append(k)
            b.append(c)
        a.append(b)
    return a

Error

score 0 · Answer 1 · answered Oct 24 '22 at 15:01

0

I try to avoid using empty lists with numba, mainly because an empty list cannot be typed. Check out nb.typeof([])

I am not sure whether your output can be preallocated but you could consider arrays. There would also be massive performance benefits. Here is an attempt:

from numba import njit, prange, int32
import numpy as np

@njit(int32[:,:,:](int32, int32), fastmath = True, parallel = True, cache = True)
def test(x, y):
    out = np.zeros((10,x,y), dtype=int32)
    for i in prange(10):
        for j in range(x):
            for k in range(y):
                out[i][j][k] = k
    return out

That said, you might indeed need lists for your application, in which case this answer might not be of much use.

answered Oct 24 '22 at 15:01

Kocas

302
1
12

Thanks for answering. I can't use numpy array as out, because of varying size of c. – Bhaskar Dhariyal Oct 24 '22 at 15:04
Understandable. My hunch is, even if you could make signatures work somehow, the performance of the whole thing might be worse than a python only implementation. if you know the max size of your elements in each level maybe you could preallocate for that and fill the rest with np.nan? – Kocas Oct 24 '22 at 15:18
1

You can initialize an empty but typed `List` with `List.empty_list()`, or alternatively use a 0-size list comprehension like `empty_f32 = [np.float32(x) for x in range(0)]`. That last one only works when inside a jitted function. – Rutger Kassies Oct 25 '22 at 09:12

score 0 · Answer 2 · answered Oct 24 '22 at 15:15

0

This worked for me.


from numba import njit, prange
from numba.typed import List

@njit(fastmath = True, parallel = True, cache = True)
def test(x, y):
    a = List()
    for i in prange(10):
        b = List()
        for j in range(4):
            c = List()
            for k in range(5):
                c.append(k)
            b.append(c)
        a.append(b)
    return a

answered Oct 24 '22 at 15:15

Bhaskar Dhariyal

1,343
2
13
31

do you have any info on how it performs? – Kocas Oct 24 '22 at 15:32

score 0 · Answer 3 · answered Oct 25 '22 at 09:40

Your signature is fine, but you need to match the type of List that you create inside the function. So a numba.typed.List instead of [].

from numba import njit, prange
from numba.typed import List
from numba.types import int32

@njit("ListType(ListType(ListType(int32)))(int32, int32)", fastmath=True, parallel=True, cache=True)
def test(x, y):
    a = List.empty_list(List.empty_list(List.empty_list(int32)))
    for i in prange(10):
        b = List.empty_list(List.empty_list(int32))
        for j in range(4):
            c = List.empty_list(int32)
            for k in range(5):
                c.append(int32(k))
            b.append(c)
        a.append(b)
    return a

I don't think you should expect much from appending to a List in parallel in this case.

How to set numba signature with nested lists?

3 Answers3