AssertionError: Unreachable condition reached (op code RERAISE executed) with batch_size > 1

import numpy as np
import numba as nb

@nb.njit(fastmath=True)
def matrix_update(feature_matrix, A_inv, temp_arrays):
    """
    JIT-compiled function that focuses on matrix inversion operations
    which likely cause the Numba error
    """
    batch_size, n_features = feature_matrix.shape
    
    # Unpack temporary arrays
    lambda_A_inv_feat_T, M, M_inv = temp_arrays
    
    # Compute A_inv @ feature_matrix.T
    for j in range(n_features):
        for i in range(batch_size):
            val = 0.0
            for k in range(n_features):
                val += A_inv[j, k] * feature_matrix[i, k]
            lambda_A_inv_feat_T[j, i] = val
    
    # Compute M = I + feature_matrix @ A_inv @ feature_matrix.T
    for i in range(batch_size):
        for j in range(batch_size):
            if i == j:
                M[i, j] = 1.0  # Identity matrix diagonal
            else:
                M[i, j] = 0.0  # Zero off-diagonal elements
    
    for i in range(batch_size):
        for j in range(batch_size):
            for k in range(n_features):
                M[i, j] += feature_matrix[i, k] * lambda_A_inv_feat_T[k, j]
    
    # Matrix inversion - this is likely where the error happens
    # Reset M_inv to zeros
    for i in range(batch_size):
        for j in range(batch_size):
            M_inv[i, j] = 0.0
    
    # Handle different batch sizes with explicit code paths
    if batch_size == 1:
        # Direct inversion for 1x1 matrices
        M_inv[0, 0] = 1.0 / M[0, 0]
    elif batch_size == 2:
        # 2x2 matrix inversion
        det = M[0, 0] * M[1, 1] - M[0, 1] * M[1, 0]
        if abs(det) > 1e-10:
            M_inv[0, 0] = M[1, 1] / det
            M_inv[0, 1] = -M[0, 1] / det
            M_inv[1, 0] = -M[1, 0] / det
            M_inv[1, 1] = M[0, 0] / det
        else:
            # Fallback to diagonal regularization
            for i in range(batch_size):
                M_inv[i, i] = 1.0 / (M[i, i] + 1e-8)
    else:
        # Simplified approach for larger matrices
        try:
            # Simple regularized diagonal approach
            for i in range(batch_size):
                M_inv[i, i] = 1.0 / (M[i, i] + 1e-8)
        except:
            # Fallback - just in case
            for i in range(batch_size):
                M_inv[i, i] = 1.0 / (M[i, i] + 1e-8)
    
    # Update A_inv using Woodbury identity (simplified)
    new_A_inv = np.zeros_like(A_inv)
    for i in range(n_features):
        for j in range(n_features):
            new_A_inv[i, j] = A_inv[i, j]
            for k in range(batch_size):
                for l in range(batch_size):
                    new_A_inv[i, j] -= lambda_A_inv_feat_T[i, k] * M_inv[k, l] * lambda_A_inv_feat_T[j, l]
    
    return new_A_inv

class MatrixUpdater:
    def __init__(self, n_features, max_batch_size=10):
        self.n_features = n_features
        self.max_batch_size = max_batch_size
        self.dtype = np.float32
        
        # Initialize model parameters
        self._A_inv = np.eye(n_features, dtype=self.dtype)
        
        # Pre-allocate temporary arrays
        self._pre_allocate_temp_arrays(max_batch_size)
        
        # Warm up JIT
        self._warm_up_jit()
    
    def _pre_allocate_temp_arrays(self, batch_size):
        """Pre-allocate temporary arrays"""
        if batch_size > self.max_batch_size:
            self.max_batch_size = batch_size
        
        # Pre-allocate temporary arrays
        self._lambda_A_inv_feat_T = np.zeros((self.n_features, self.max_batch_size), dtype=self.dtype)
        self._M = np.zeros((self.max_batch_size, self.max_batch_size), dtype=self.dtype)
        self._M_inv = np.zeros((self.max_batch_size, self.max_batch_size), dtype=self.dtype)
        
        # Package arrays for Numba function
        self._temp_arrays = (self._lambda_A_inv_feat_T, self._M, self._M_inv)
    
    def _warm_up_jit(self):
        """Warm up the JIT compiler"""
        try:
            # Create tiny dummy data
            X_dummy = np.ones((1, self.n_features), dtype=self.dtype)
            
            # Create small temporary arrays for warm-up
            lambda_A_inv_feat_T = np.zeros((self.n_features, 1), dtype=self.dtype)
            M = np.zeros((1, 1), dtype=self.dtype)
            M_inv = np.zeros((1, 1), dtype=self.dtype)
            
            # Run JIT function to compile it
            _ = matrix_update(X_dummy, self._A_inv, (lambda_A_inv_feat_T, M, M_inv))
            print("JIT compilation warmed up for matrix operations")
        except Exception as e:
            print(f"JIT warm-up failed: {e}")
    
    def update(self, X):
        """Update the matrix with new data"""
        if isinstance(X, list):
            X = np.array(X, dtype=self.dtype)
        elif X.dtype != self.dtype:
            X = X.astype(self.dtype)
        
        # Ensure correct shape
        if len(X.shape) == 1:
            X = X.reshape(1, -1)
        
        batch_size = X.shape[0]
        
        # Ensure temporary arrays are large enough
        if batch_size > self.max_batch_size:
            print(f"Resizing arrays for batch size {batch_size}")
            self._pre_allocate_temp_arrays(batch_size)
        
        # Perform the update operation
        self._A_inv = matrix_update(X, self._A_inv, self._temp_arrays)
        
        return self._A_inv

def run_test():
    import time
    
    # Create sample data
    np.random.seed(42)
    n_samples = 1000
    n_features = 5
    
    # Generate data
    X = np.random.randn(n_samples, n_features).astype(np.float32)
    
    # Test with different batch sizes
    for batch_size in [1, 10, 100]:
        print(f"\nTesting with batch_size = {batch_size}")
        
        # Initialize updater
        updater = MatrixUpdater(n_features=n_features, max_batch_size=batch_size)
        
        # Track time
        start_time = time.time()
        
        # Update in batches
        for i in range(0, n_samples, batch_size):
            end_idx = min(i + batch_size, n_samples)
            X_batch = X[i:end_idx]
            A_inv = updater.update(X_batch)
        
        # Print results
        elapsed = time.time() - start_time
        print(f"Elapsed time: {elapsed:.6f} seconds")
        print(f"Final A_inv shape: {A_inv.shape}")

if __name__ == "__main__":
    run_test()

python3 -m test_numba

Testing with batch_size = 1
JIT compilation warmed up for matrix operations
Elapsed time: 0.001606 seconds
Final A_inv shape: (5, 5)

Testing with batch_size = 10
JIT compilation warmed up for matrix operations
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "test_numba.py", line 181, in <module>
    run_test()
  File "test_numba.py", line 173, in run_test
    A_inv = updater.update(X_batch)
            ^^^^^^^^^^^^^^^^^^^^^^^
  File "test_numba.py", line 144, in update
    self._A_inv = matrix_update(X, self._A_inv, self._temp_arrays)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "test_numba.py", line 66, in matrix_update
    for i in range(batch_size):
AssertionError: Unreachable condition reached (op code RERAISE executed)
-------------------------------------------------------------------------------
This should not have happened, a problem has occurred in Numba's internals.
You are currently using Numba version 0.61.0.

Please report the error message and traceback, along with a minimal reproducer
at: https://github.com/numba/numba/issues/new?template=bug_report.md

If more help is needed please feel free to speak to the Numba core developers
directly at: https://gitter.im/numba/numba

Thanks in advance for your help in improving Numba!