EDIT 2.
It turns out the problem below is due to a bug in numba (version 0.54.1) which is now being examined by the Numba dev team. Here is the corresponding GH issue.
========
I have a very simple Python/Numpy function and I want to compile it with Numba.
Everything works great for standard compilation, but if I use parallel=True then suddenly I get nonsensical corrupted data and I do not have a clue why.
Below is a reproducible example.
import numpy as np
from numba import njit, prange
@njit(boundscheck=False, nogil=True, parallel=False)
def ubcm_jacobian(
D: np.ndarray,
b: np.ndarray
) -> np.ndarray:
"""Jacobian of UBCM negative log-likelihood.
Parameters
----------
D
Degree sequence array.
b
Model parameters.
"""
B = np.empty_like(b)
b = np.exp(b)
for i in prange(len(b)):
bi = b[i]
x = bi * b
B[i] = (x / (1 + x)).sum()
b2 = b**2
B -= b2 / (1 + b2)
return B - D
Above is the function and below is simple data for which the error should reproduce.
D = np.arange(1, 11)
b = np.log(D)
ubcm_jacobian(D, b).sum() # correct result: 28.9899 (roughly)
But if I compile the above function with parallel=True then I get nonsensical results, basically as if B array was not populated with values within the main loop.
And this is surprising because the whole routine seems to be trivially parallelizable. So what am I doing wrong?
====== Below I add the output of parallel diagnostics produced by Numba (it seems to be fine?)
================================================================================
Parallel Accelerator Optimizing: Function ubcm_jacobian, <ipython-
input-128-8dd321ee9df4> (1)
================================================================================
Parallel loop listing for Function ubcm_jacobian, <ipython-input-128-8dd321ee9df4> (1)
----------------------------------------------------------|loop #ID
@njit(boundscheck=False, nogil=True, parallel=True) |
def ubcm_jacobian( |
D: np.ndarray, |
b: np.ndarray |
) -> np.ndarray: |
"""Jacobian of UBCM negative log-likelihood. |
|
Parameters |
---------- |
D |
Degree sequence array. |
b |
Model parameters. |
""" |
B = np.empty_like(b) |
b = np.exp(b)-----------------------------------------| #39
|
for i in prange(len(b)):------------------------------| #46
bi = b[i] |
x = bi * b-------------------------------------| #40
B[i] = (x / (1 + x)).sum()------------------------| #41, 45
|
b2 = b**2---------------------------------------------| #42
B -= b2 / (1 + b2)------------------------------------| #43, 47
|
return B - D------------------------------------------| #44
------------------------------ After Optimisation ------------------------------
Parallel region 0:
+--39 (parallel)
+--40 (serial, fused with loop(s): 41, 45)
Parallel region 0 (loop #39) had 2 loop(s) fused and 1 loop(s) serialized as
part of the larger parallel loop (#39).
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
EDIT
After updating to the latest numba version (0.54.1) the changed but it is still not correct. Now with parallel=True I get a result which is much closer to the correct one given by the function compiled with parallel=False, but it is still not correct. More concretely, now I get roughly 31.99498 but the correct answer is 28.9899.
So still, something is not working.