How to Solve Numba Lowering error?
<p>I have a function, which I am trying to speed up using the @jit decorator from Numba module. For me it is essential to speed this up as much as possible, because my main code calls upon this function for millions of times. Here is my function:</p>

<pre><code>from numba import jit, types
import Sweep    #My own module, works fine


@jit(types.Tuple((types.complex128[:], types.float64[:]))(types.complex128[:], types.complex128[:], types.float64[:], types.float64[:], types.float64))
def MultiModeSL(Ef, Ef2, Nf, u, tijd ):
  dEdt= np.zeros(nrModes, dtype=np.complex128)
  dNdt0= np.zeros(nrMoments, dtype=np.complex128)
  Efcon = np.conjugate(Ef)

  for j in range(nrModes):
    for n in range(nrMoments):
      dEdt +=  0.5 * CMx[:,j,n,0] * dg * (1+ A*1j) * Nf[n] * Ef[j] * np.exp( 1j* (Sweep.omega[j]-Sweep.omega) *tijd)
      for k in range(nrModes):
        if n==0:
          dNdt0 += g* CMx[j, k, 0,:] * Efcon[j] * Ef[k] * np.exp( 1j* (Sweep.omega[k]-Sweep.omega[j]) *tijd) 
        dNdt0 += dg*(1+A*1j) * CMx[j,k,n,:] * Nf[n] * Efcon[j] * Ef[k] * np.exp( 1j* (Sweep.omega[k]-Sweep.omega[j]) *tijd) 

  dEdt += - 0.5*(pd-g)*Ef +     fbr*Ef2  + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) ) 
  dNdt = Sweep.Jn - Nf*ed - dNdt0.real
  return dEdt, dNdt
</code></pre>

<p>The function works perfectly well, without the Jit decorator. However, when I run it with the @jit, I get this error:</p>

<pre><code>numba.errors.LoweringError: Failed at object (object mode frontend)
Failed at object (object mode backend)
dEdt.1
File "Functions.py", line 82
[1] During: lowering "$237 = call $236(Ef, Ef2, Efcon, Nf, dEdt.1, dNdt0, tijd, u)" at /home/humblebee/MEGA/GUI RC/General_Formula/Functions.py (82)
</code></pre>

<p>Line 82 corresponds to the For loop with j as iterator. </p>

<p>Can you help me out?</p>

<p><strong><em>EDIT:</em></strong>
Based on Peter's suggestion and combining it with Einsum, I was able to remove the loops. This made my function <strong>3</strong> times faster.  Here is the new code:</p>

<pre><code>def MultiModeSL(Ef, Ef2, Nf, u, tijd ):
  dEdt= np.zeros(nrModes, dtype=np.complex128)
  dNdt0= np.zeros(nrMoments, dtype=np.complex128)
  Efcon = np.conjugate(Ef)
  dEdt = 0.5*  np.einsum("k, jkm, mk, kj -&gt; j",  dg*(1+A*1j), CMx[:, :, :, 0],  (Ef[:] * Nf[:, None] ),  np.exp( 1j* (OMEGA[:, None]-OMEGA) *tijd))
  dEdt += - 0.5*(pd-g)*Ef + fbr*Ef2  + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) )

  dNdt = - np.einsum("j, jkm, jk, kj ", g, CMx[:,:,:,0], (Ef*Efcon[:,None]),  np.exp( 1j* (OMEGA[:, None]-OMEGA) *tijd))
  dNdt += -np.einsum("j, j, jknm, kjm, kj",dg, (1+A*1j), CMx, (Nf[:]*Efcon[:,None]*Ef[:,None,None]), np.exp( 1j* (OMEGA[:, None]-OMEGA) *tijd)  )
  dNdt += JN - Nf*ed
  return dNdt
</code></pre>

<p>Can you suggest more techniques to speed this up?</p>

<p>I can't see from your code why this isn't vectorizable.  Vectorizing can speed up this kind of Python code by around 100x.  Not sure of how it does relative to jit.</p>

<p>It looks like you could, for instance, take your dEdt out of the loop, and compute it in one step with something like :</p>

<pre><code>dEdt = 0.5 * (Cmx[:, :, :, 0] * dg * (1+A*1j) * Nf[:] * Ef[:, None] * np.exp( 1j* (Sweep.omega[None, :, None, None]-Sweep.omega) *tijd)).sum(axis=2).sum(axis=1) - 0.5*(pd-g)*Ef + fbr*Ef2  + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) )
</code></pre>

<p>(Though I don't really know what the dimensionality of your Sweet.omega is).</p>

python numpy numba