# Another List of Arrays question (Numpy array inside a List Comprehension)

**URL:** <https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804>\
**Category:** Support: How do I do ...?\
**Created:** [February 28, 2023, 3:41am UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804 "2023-02-28T03:41:14Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![andressommerhoff](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/andressommerhoff/32/788_2.png) [@andressommerhoff](https://numba.discourse.group/u/andressommerhoff)\
**Post date:** [February 28, 2023, 3:41am UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/1 "2023-02-28T03:41:14Z")

</div>

Hi team!!

Someone knows how to make List Comprehension work with NumPy arrays inside?

I’m having trouble getting the following jitted function to work, which is a simplified version of a larger code:

```auto
@nb.njit()
def multimask(search_arr: np.ndarray, bigarray: np.ndarray):
      return np.array([bigarray == q for q in search_arr]) # <-Numba error

```

It doesn’t work when jitted. Below are more explanations and questions…

* * *

## Expected return:

The function should return a boolean 2D array where each row shows where the corresponding value in `search_arr` was found in `bigarray` . For example:

```auto
>>> a = np.array([10,20,10,30,10])
>>> s = np.array([10,30])
>>> multimask(s, a)
array([[True, False, True, False, True], # Found 10 in 1st, 3rd, 5th
       [False, False, False, True, False]]) # Found 30 in 4th item

```

But I get the following error:

> **.... click here to see error ...**
>
> ```auto
> ERROR:
> No implementation of function Function(<built-in function setitem>) found for signature:
>  
> >>> setitem(array(undefined, 1d, C), int64, array(bool, 1d, C))
>  
> There are 16 candidate implementations:
> - Of which 16 did not match due to:
> Overload of function 'setitem': File: <numerous>: Line N/A.
> With argument(s): '(array(undefined, 1d, C), int64, array(bool, 1d, C))':
> No match.
> 
> ```

* * *

## Work around, but verbose…

While I managed to work around the jit error by opening the List Comprehension into a explicit loop, I don’t like the verbosity:

```auto
@nb.njit()
def multimask_okcompile(search_arr:np.ndarray, bigarray:np.ndarray):
    size_search = len(search_arr)
    size_bigarray = len(bigarray)
    arr_mask = np.empty((size_search, size_bigarray), dtype="bool")
    for pos in range(size_search):
        q = search_arr[pos]
        arr_mask[pos,:] = (bigarray == q)
    return arr_mask

```

… I would prefer a cleaner pythonic solution and insist on List Comprehension or similar (if possible in Numba or Numpy).

* * *

## Any ideas?

I know that list comprehension of simple types works fine in Numba, but I’m not sure **how to make List Comprehension work with NumPy arrays inside**. I’m wondering:

- Do I need to add (the right) **type hints**?
- Can **TypedList** be used somehow to avoid explicit loop?
- or maybe I need to **overloading** `np.array` or the built-in function `setitem` in some way? (I have taken a look to Numba issue 4470 ([“Can’t create a numpy array from a numpy array”](https://github.com/numba/numba/issues/4470)) to overload “np.array” to receive arrays. Fine for me if something like this works while waiting for Numba support arrays inside list comprehension)

Any direction or help on this will be highly appreciated!

* * *

## What have I explored so far:

Before posting here, I made a lot of code testing and investigation (in the documentation, this forum, and even ChatGPT 😉). So far, I have tried:

- Adding type hints within the function in different ways.
- Defining the types outside the JIT scope and adding them as a `dtype` parameter.
- Searching unsuccessfully for a different NumPy function to replace the Numba code (I have tried `np.isin` , `np.vstack` , and different NumPy array creation functions) in case it is not possible to make list comprehension work with NumPy arrays inside.
- Reading a lot about Numba arrays of arrays or lists of arrays (for example, [“passing a list of numpy arrays into np array with numba”](https://numba.discourse.group/t/passing-a-list-of-numpy-arrays-into-np-array-with-numba/278)). However, something like **make\_2d()** is just similar to my verbose workaround.

---

<div class="post-metadata">

**Author:** ![Rutger](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/rutger/32/446_2.png) [@Rutger](https://numba.discourse.group/u/Rutger)\
**Post date:** [February 28, 2023, 7:50am UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/2 "2023-02-28T07:50:51Z")

</div>

You could consider using `guvectorize`, that simplifies it to a 1D problem, which is perhaps less verbose to you.

```python
@guvectorize("(n),()->(n)")
def multimask(a, q, out):
    for i in range(a.size):
        out[i] = a[i] == q
        
a = np.array([10,20,10,30,10])
s = np.array([10,30])

out = np.empty((s.size, a.size), dtype="bool")
multimask(a, s, out)

```

You can have Numba do the output allocation if you’re capable of pinning down the datatypes it should use, for example:

```python
@guvectorize("void(int32[:], int32, boolean[:])", "(n),()->(n)")
def multimask(a, q, out):
    for i in range(a.size):
        out[i] = a[i] == q

multimask(a, s)

```

You can flip the “axis” of vectorization if needed. If the sizes between the two arrays are very different, it might matter for performance.

---

<div class="post-metadata">

**Author:** ![andressommerhoff](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/andressommerhoff/32/788_2.png) [@andressommerhoff](https://numba.discourse.group/u/andressommerhoff)\
**Post date:** [February 28, 2023, 3:12pm UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/3 "2023-02-28T15:12:40Z")

</div>

Thanks, @Rutger, for introducing me to `guvectorize` . I plan to use it in the future. However, it would be great if I could keep the NumPy arrays inside a Comprehension List in Numba. Is there a known way to make Numba handle this? Perhaps it can be implemented by overloading something?

Why to insist in _arrays_ inside _Comprehension List_?

- The original Python code can be Jitted **without modification**. This is similar to how the solution to Numba [issue #4470](https://github.com/numba/numba/issues/4470) opens the door to directly use np.array with arrays in Numba, avoiding awkward mangling modifications before Jitting the code.
- It enhances code **clarity** and **expressiveness**.
- It seems to be a super **powerful tool** if we have as an alternative in Numba.

I’m taking this as a personal challenge and I’m willing to try a few more things before giving up. Do you have any ideas for me to try?

---

<div class="post-metadata">

**Author:** ![andressommerhoff](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/andressommerhoff/32/788_2.png) [@andressommerhoff](https://numba.discourse.group/u/andressommerhoff)\
**Post date:** [February 28, 2023, 3:21pm UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/4 "2023-02-28T15:21:10Z")

</div>

For a first attempt at my challenge, I’d like to explore a solution similar to the one presented in [issue #4470](https://github.com/numba/numba/issues/4470) .

I may be mistaken, but I have a hunch that overloading the `setitem` function or the list creation method (if there is one in Numba) could be the way to go to implement comprehension lists with arrays inside. The documentation for `@overload` and `@overload_method` are clear, but I’m unsure which object should be overloaded for the comprehension list case.

If anyone have information or guidance on this topic, I would greatly appreciate it!

_PS: To make the quest less generic, I am willing to focus on one specific type of Numpy array. For example, it could be a 1D boolean array inside a comprehension list iterating over integers._

---

<div class="post-metadata">

**Author:** ![Rutger](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/rutger/32/446_2.png) [@Rutger](https://numba.discourse.group/u/Rutger)\
**Post date:** [March 1, 2023, 10:50am UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/5 "2023-03-01T10:50:21Z")

</div>

Thanks for elaborating on why you want to use the comprehensions.

- It’s indeed a nice feature if you can have “regular” Numpy code that can be jitted without modifications, that allows seamlessly switching between both without needing to maintain redundant implementations. I do find though that Numba especially shines when abandoning the Numpy-vectorized notation (eg explicit loops etc), but it depends on the specific case of course.
- I disagree somewhat with calling comprehensions clear. I use them all the time, it’s a very powerful notation, but not the most readable to me personally. This is very subjective of course.
- It would indeed be great to have Numba support this.

Specifically regarding your requirement of wanting to keep the code generic Numpy, an alternative way of achieving that with the current release of Numba is to utilize the broadcasting mechanism.

I would normally write this myself as `a[None,:] == s[:,None]` or equivalently `a[np.newaxis,:] == s[:,np.newaxis]`, both of which aren’t supported at the moment. But the same result is obtained when using the reshape method, which does work:

```python
@njit
def multimask(a, s):
    return a.reshape(1,-1) == s.reshape(-1,1)

multimask(a,s)

```

It shows good performance for me, slightly slower compared to the above mentioned `guvectorized` alternative, but faster compared to the same un-jitted function (pure Numpy).

---

<div class="post-metadata">

**Author:** ![andressommerhoff](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/andressommerhoff/32/788_2.png) [@andressommerhoff](https://numba.discourse.group/u/andressommerhoff)\
**Post date:** [March 6, 2023, 7:51pm UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/6 "2023-03-06T19:51:43Z")

</div>

Thanks @Rutger! For my original problem your last solution is what I needed! Super clean `a.reshape(1,-1) == s.reshape(-1,1)` compared with my first comprehension list (Looking forward to the day when Numba will support its equivalent with an even cleaner syntax: `a[None,:] == s[:,None]`).

> [@Rutger](#):
>
> I disagree somewhat with calling comprehensions clear. I use them all the time, it’s a very powerful notation, but not the most readable to me personally. This is very subjective of course.

👆🏽 Of course, it is subjective, but I would say that it depends on the case. Having more than two `for` loops inside a list comprehension is too cryptic for my taste. However, a simple list comprehension with one or two for loops plus an if statement is usually easy to grasp than the explicit loop with an if condition… in my subjective opinion…

> [@Rutger](#):
>
> It would indeed be great to have Numba support this.

👆🏽 I fully agree. I’m looking forward to seeing that implemented someday.

Thank you again Rutger for all your help!

PS: no news unfortunately from my side regarding using @overload for making the simple use case: a 1D boolean array inside a comprehension list iterating over integers. (But I haven’t been able to make many attempts either).

---

<div class="post-metadata">

**Author:** ![sschaer](https://avatars.discourse-cdn.com/v4/letter/s/5daacb/32.png) [@sschaer](https://numba.discourse.group/u/sschaer)\
**Post date:** [March 6, 2023, 8:58pm UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/7 "2023-03-06T20:58:05Z")

</div>

Hi @andressommerhoff

The solution of @Rutger is indeed very elegant. As for the slowdown compared to the `guvectorize` function, I can give you the tip to only reshape one of the arguments. E.g

```auto
@njit
def multimask(a, s):
    return a == s.reshape(-1,1)

```

will be considerably faster and even faster than using `guvectorize`.

---

<div class="post-metadata">

**Author:** ![andressommerhoff](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/andressommerhoff/32/788_2.png) [@andressommerhoff](https://numba.discourse.group/u/andressommerhoff)\
**Post date:** [March 8, 2023, 1:54pm UTC](https://numba.discourse.group/t/another-list-of-arrays-question-numpy-array-inside-a-list-comprehension/1804/8 "2023-03-08T13:54:30Z")

</div>

> [@sschaer](#):
>
> ```auto
> @njit
> def multimask(a, s):
> return a == s.reshape(-1,1)
> 
> ```

Such a nice solution @sschaer! You nailed it! (And now I feel more like a newbie than I thought I was, hahaha)