I am trying to use numba.jit to compile a function that takes several inputs and returns one Tuple. I specified the types of inputs and output within the jit decoration, thinking that this would have help. However, I get a warning saying that a list containing a tuple cannot be converted into a tuple and the numba.jit falls back into object-mode.
I am building a script to model a physical 2D phenomenon. Here below there is the code of the function I am trying to compile with numba.jit (a bit simplified). There are some parameters that are not listed in the input. These are constant parameters, defined previously, of type float with size 1. The critical thing here is that I declared that I will receive as an output a type Tuple(array(float64, 3d, C), array(float64, 2d, A)).
@nb.jit(nb.types.containers.Tuple((nb.types.float64[:,:,::1],nb.types.float64[:,:]))(nb.types.float64,nb.types.containers.Tuple((nb.types.float64[:,:,:],nb.types.float64[:,:])),nb.types.float64[:,:,:],nb.types.float64[:,:,:],nb.types.float64[:,:],nb.types.float64[:,:,:],nb.types.float64[:,:]))
def model_f(t,dvar,dNdy,N_y,N_x,r,coeff2):
X = dvar[0]
X_ch = dvar[1]
# trivial operations on matrix element-wise:
i = 0
for row in X[0,:,0]:
j = 0
for col in X[0,0,:]:
N_y[:,i,j] = DGM_N(X,i,j)
[r[:,i,j]] = RRates(X,T,i,j)
j += 1
i += 1
# gradient of N_y:
comp = 0
for c_plane in X[:,0,0]:
i = 0
for row in X[0,:,0]:
if i == 0:
dNdy[comp,i,:] = (N_y[comp,i+1,:]-N_y[comp,i,:])/dy
elif i == NY-1:
dNdy[comp,i,:] = (N_y[comp,i,:]-N_y[comp,i-1,:])/dy
else:
dNdy[comp,i,:] = (N_y[comp,i+1,:]-N_y[comp,i-1,:])/(2*dy)
i += 1
comp += 1
coeff1 = 1/(-dNdy+r)
N_x = speed*X_ch
i = 0
for cell in x_axis:
if i == 0:
coeff2[:,i] = (-N_x[:,i]+No)/dx-N_y[:,0,i]/channel_height
else:
coeff2[:,i] = (-N_x[:,i]+N_x[:,i-1])/dx-N_y[:,0,i]/channel_height
i += 1
return [(coeff1,coeff2)]
But, when I try to compile the function, I get the following error:
NumbaWarning:
Compilation is falling back to object mode WITH looplifting enabled because Function "model_f" failed type inference due to: No conversion from list(Tuple(array(float64, 3d, C), array(float64, 2d, A)))<iv=None> to Tuple(array(float64, 3d, C), array(float64, 2d, A)) for '$1656return_value.4', defined at None
File "soe_model_v2.1.py", line 338:
def model_f(t,dvar,T,mu_fuel,mu_air,dXdy,dNdy,N_y,N_x,r,Dkn,Deffbin,coeff2,result):
<source elided>
return [(coeff1,coeff2)]
^
During: typing of assignment at c:\users\matfe\...\soe_model_v2.1.py (338)
File "soe_model_v2.1.py", line 338:
def model_f(t,dvar,T,mu_fuel,mu_air,dXdy,dNdy,N_y,N_x,r,Dkn,Deffbin,coeff2,result):
<source elided>
return [(coeff1,coeff2)]
^
@nb.jit(nb.types.containers.Tuple((nb.types.float64[:,:,::1],nb.types.float64[:,:]))(nb.types.float64,nb.types.containers.Tuple((nb.types.float64[:,:,:],nb.types.float64[:,:])),nb.types.float64,nb.types.float64,nb.types.float64,nb.types.float64[:,:,:],nb.types.float64[:,:,:],nb.types.float64[:,:,:],nb.types.float64[:,:],nb.types.float64[:,:,:],nb.types.float64[:],nb.types.float64[:,:],nb.types.float64[:,:],nb.types.containers.Tuple((nb.types.float64[:,:,:],nb.types.float64[:,:]))))
You can see that in the original function code I have more inputs to my function, that are also inferred during the jit decoration (which makes the error message quite long) but that should not matter here.
So it seems like my "return" object is seen as a list containing a tuple, instead of just being a tuple. Is this the reason why it does not compile? Does anybody have any suggestions?
Returning multiple values in Python already defaults to a tuple, so you should be able to simplify it to something like the toy example below:
import numba as nb
from numba.types import Tuple, float64, int64
import numpy as np
@nb.njit(Tuple((float64[:,:,::1], float64[:,:]))(int64))
def model_f(x):
a = np.random.randn(x, x, x)
b = np.random.randn(x, x)
return a, b
model_f(2)
This will compile correctly, and return a tuple with the two arrays.