What is the correct numba decorator to return an array of different size than the parameter passed?

Question

I want to use numba to vectorize a function to count the occurrences of unique values. The function receives a numpy array of any length, and returns a numpy array of length 257.

But I do not understand how to specify the decorator. It doesn't' even compile. I read the documentation here, but i get empty handed. It says nothing about different array sizes.

@nb.guvectorize([(nb.uint8[:], nb.uint64[:])], "(n) -> (m)", target="parallel")
def count_occurrences(byte_view):
    """
    Counts the occurrences of each element in a byte array and returns a new array with the counts.
    """
    # adds a zero to the beginning of the array, for convenience
    count = np.zeros(1 + 256, dtype=np.uint64)
    count[1 + byte_view] += 1
    return count


sample = np.random.randint(1, 100, 100, dtype=np.uint8)
counts = count_occurrences(sample)

Answer 1

If both are 1D arrays, there isn't anything that can be done in parallel with this setup. And then there doesn't seem to be much benefit of using guvectorize over njit in this case. Especially because njit would allow returning an array with a different shape than any of the inputs.

It's not possible to return an array of unknown shape with guvectorize, so you'll have to provide that as an input argument as well. And you should never return anything explicitly from a guvectorize function to begin with. That might be more intuitive if you consider that Numba (ideally) would parallelize the execution of the function you provide, meaning that in this case count would be initialized multiple times instead of only once (upfront).

It's not clear what the indexing with byte_view means in this context, if it contains the indices that need to be indexed, you can just loop over them as shown below.

import numpy as np
import numba as nb

@nb.guvectorize("void(uint8[:], uint64[:])", "(n),(m)", target="cpu")
def count_occurrences(byte_view, count):
    """
    Counts the occurrences of each element in a byte array and returns a new array with the counts.
    """
    for idx in byte_view: 
        count[1 + idx] += 1


sample = np.random.randint(1, 100, 100, dtype=np.uint8)

# adds a zero to the beginning of the array, for convenience
counts = np.zeros(1 + 256, dtype=np.uint64)

count_occurrences(sample, counts)

The snippet above makes it work, but as said, it doesn't add much over njit, and doesn't utilize what guvectorize normally adds over njit.