Numba "LLVM IR parsing error" seemingly because numba.complex128[:,:].shape has a default type? - numpy

Question

Please bear with the long question.
Numba encounters "LLVM IR parsing error" in my code seemingly due to defualt typing of np.complex128.shape, but I could not find any documentations saying that np.complex128.shape, numba.complex128.shape or prange have default types.
Minimal workable reproduction:
import numpy as np
from numba import jit, njit, prange
from numba import complex128, int32 # import jit value types
# Invert an (n,n) submatrix of a (m>n,n) rectangular matrix by taking the first
# n rows. "Taking the first n rows" is motivated by the RHS being rank n.
#
# -- Input --
# (m,n) matrix A
#
# -- Return --
# (m,m) matrix A_inv
#njit(complex128[:,:](complex128[:,:]))
def inv_square_jit(in_matrix):
if in_matrix.ndim != 2:
raise ValueError("Input should be 2d array")
n_row = in_matrix.shape[0]
n_col = in_matrix.shape[1]
if n_row<=n_col:
raise ValueError("Input should have more rows than cols")
# Remove specfied column (slightly faster than delete)
# and remove extra rows
sqinv = np.linalg.inv(in_matrix[:n_col, :])
padded = np.zeros((n_row, n_row), dtype = np.complex128)
padded[:len(sqinv), :len(sqinv)] = sqinv
return(padded)
# Solve degenerate underdetermined equation system
# -- Input --
# (m,n+1), (m,n), rank-n 2d np arrays A, B
# n-dim np array-like vb
# or
# (m,n+1), rank n+1 A,
# m-dim np array-like v_rhs
#
# vb can be any array-like item, and is not necessarily 1d.
# Implemented with ChiPhiFunc in mind.
#
# -- Return --
# n+1 np array-like va
#
# -- Note --
# For recursion relations with ChiPhiFunc's, A and B should come from
# convolution matrices. That still needs implementation.
#njit(complex128[:](complex128[:,:], complex128[:]))
def solve_degenerate_jit(A, v_rhs):
n_dim = A.shape[1]
if A.shape[0] != v_rhs.shape[0]:
raise ValueError("solve_underdetermined: A, v_rhs must have the same number of rows")
A_inv = np.ascontiguousarray(inv_square_jit(A))
# This vector is actually m-dim, with m-n blank elems at the end.
va = (A_inv#np.ascontiguousarray(v_rhs))[:n_dim]
return(va)
# #njit(complex128[:](complex128[:,:], complex128[:,:], complex128[:]))
# def solve_degenerate_jit(A, B, vb):
# B_cont = np.ascontiguousarray(B)
# vb_cont = np.ascontiguousarray(vb)
# return(solve_degenerate_jit(A, B_cont#vb_cont))
# Generate convolution operator from a for an n_dim vector.
#njit(complex128[:,:](complex128[:], int32))
def conv_matrix(vec, n_dim):
out_transposed = np.zeros((n_dim,len(vec)+n_dim-1), dtype = np.complex128)
for i in prange(n_dim):
out_transposed[i, i:i+len(vec)] = vec
return(out_transposed.T)
# For solving a*va = v_rhs, where va, vb have the same number of dimensions.
# In the context below, "#dim" represents number of chi mode components.
#
# -- Input --
# v_source_A: 2d matrix, content of ChiPhiFuncGrid, #dim = a
# v_rhs: 2d matrix, content of ChiPhiFuncGrid, #dim = m
# rank_rhs: int, rank of v_rhs (and correct answer)
# -- Output --
# va: 2d matrix, content of ChiPhiFuncGrid. Has #dim = rank_rhs
#njit(complex128[:,:](complex128[:,:], complex128[:,:], int32), parallel=True)
def batch_degen_jit(v_source_A, v_rhs, rank_rhs):
# if type(v_source_A) is not ChiPhiFuncGrid or type(v_source_B) is not ChiPhiFuncGrid:
# raise TypeError('batch_underdetermined_deconv: input should be ChiPhiFuncGrid.')
A_slices = np.ascontiguousarray(v_source_A.T) # now the axis 0 is phi grid
v_rhs_slices = np.ascontiguousarray(v_rhs.T) # now the axis 0 is phi grid
# axis 0 is phi grid, axis 1 is chi mode
va_transposed = np.zeros((len(A_slices), rank_rhs), dtype = np.complex128)
if len(A_slices) != len(v_rhs_slices):
raise ValueError('batch_underdetermined_deconv: A, v_rhs must have the same number of phi grids.')
if len(v_source_A) + rank_rhs - 1 != len(v_rhs):
raise ValueError('batch_underdetermined_deconv: #dim_A + rank_rhs - 1 = #dim_v_rhs must hold.')
for i in prange(len(A_slices)):
A_conv_matrix_i = conv_matrix(A_slices[i], rank_rhs)
# ********** Removing this line somehow makes it compile **********
va_transposed[i, :] = solve_degenerate_jit(A_conv_matrix_i,v_rhs_slices[i])
# ********** Removing this line somehow makes it compile **********
return va_transposed.T
The code compiles fine with parallel=False for the last method. However, with parallel=True, error occurs in for i in prange(len(A_slices)): of def batch_degen_jit(v_source_A, v_rhs, rank_rhs):, seemingly because solve_degenerate_jit(complex128[:,:], complex128[:], int32) accepts int32, but the prange(len(A_slices)) produces int64. Replacing all int32 with int64 solves the problem. Removing the *-marked line also makes it compile.
Error:
LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython mode backend)
LLVM IR parsing error
<string>:1278:34: error: '%.777' defined with type 'i64' but expected 'i32'
%".778" = icmp eq i32 %".776", %".777"
^
File "<ipython-input-24-fa65c2d527fa>", line 104:
def batch_degen_jit(v_source_A, v_rhs, rank_rhs):
<source elided>
raise ValueError('batch_underdetermined_deconv: #dim_A + rank_rhs - 1 = #dim_v_rhs must hold.')
for i in prange(len(A_slices)):
^
During: lowering "id=17[LoopNest(index_variable = parfor_index.1805, range = (0, $10call_method.4_size0.1767, 1))]{120: <ir.Block at <ipython-input-24-fa65c2d527fa> (104)>}Var(parfor_index.1805, <ipython-input-24-fa65c2d527fa>:104)" at <ipython-input-24-fa65c2d527fa> (104)
Why is this the case?
Thank you!
(P.S. here's a test case for the methods:
convolver = np.random.rand(10,3)
correct_answer = np.random.rand(10,5)
rhs = np.zeros((10,7))
for i in range(10):
rhs[i] = np.convolve(convolver[i], correct_answer[i])
print(batch_degen_jit(np.complex128(convolver).T, np.complex128(rhs).T, 5))
)