Numba prange() Inception - laser-base/laser-core GitHub Wiki

Does using numba.prange() inside a numba.prange() loop interfere with random number generation?

Experience says it does - a version of the polio model did this and the PRNG repeatability test started failing. The test passed again when the inner prange() was converted back to simple range().

However, the following test code passes on an 8-core GitHub Codespace (inception.py):

import laser_core.random as lcr
import numba as nb
import numpy as np

prng = lcr.seed(20250703)

NUM_PEOPLE = 1 << 23

n_threads = nb.get_num_threads()
draws = prng.random(n_threads)
sum = draws.sum()
counts = np.round(draws / sum * NUM_PEOPLE).astype(np.int32)
counts[-1] = NUM_PEOPLE - counts[:-1].sum()  # Ensure total is exactly NUM_PEOPLE
starts = np.zeros(n_threads, dtype=np.int32)
starts[1:] = counts.cumsum()[:-1]

@nb.njit((nb.int32[:], nb.int32[:], nb.int32[:], nb.float32[:]), parallel=True)
def process_people(counts, starts, cores, values):
    for thread in nb.prange(nb.get_num_threads()):
        start = starts[thread]
        for person in nb.prange(counts[thread]):
            idx = start + person
            cores[idx] = nb.get_thread_id()
            values[idx] = np.random.beta(2, 5)

    return

cores1 = np.zeros(NUM_PEOPLE, dtype=np.int32)
values1 = np.zeros(NUM_PEOPLE, dtype=np.float32)
lcr.seed(20250703)
process_people(counts, starts, cores1, values1)

cores2 = np.zeros(NUM_PEOPLE, dtype=np.int32)
values2 = np.zeros(NUM_PEOPLE, dtype=np.float32)
lcr.seed(20250703)
process_people(counts, starts, cores2, values2)

print(f"{np.all(cores2  == cores1 )=}")
print(f"{np.all(values2 == values1)=}")

print("done")
$ python inception.py 
np.all(cores2  == cores1 )=True
np.all(values2 == values1)=True
done

Based on the values in cores1 (and cores2), the inner numba.prange() appears to be ignored and converted to a simple range().

Package versions:

  • numpy==1.26.4
  • numba==0.61.2
  • llvmlite==0.44.0

Also works with the latest Numba supported NumPy:

  • numpy==2.2.6
  • numba==0.61.2
  • llvmlite==0.44.0

Also works on Apple Silicon Macbook with 14 cores and the package versions immediately above (numpy==2.2.6).