How can I rewrite the indexing so the function compiles?

[sschaer]

Hi @ofk123

I think you are right about the origin of the error. Fortunately, there are many ways to make this work. It is just a compromise between readability, brevity and efficiency.
For example, you can always write it with explicit loops. This is also more efficient in your case, because you only have to go through the arrays once.

Here you find two ways to implement this. Hope this helps.

import numpy as np
import numba as nb

n = 1_000
a = np.arange(n**2, dtype=float).reshape(n, n)
b = np.arange(n**2, dtype=float).reshape(n, n)*10

# @nb.njit
def func_jit1(a, b):
    dist = a - b
    condition = dist < 0.0
    dist[condition] = dist[condition] + 100.0
    return dist

@nb.njit(parallel=False)
def func_jit2(a, b):
    dist = np.empty_like(a)
    for i in nb.prange(a.shape[0]):
        for j in range(a.shape[1]):
            dist[i, j] = a[i, j] - b[i, j]
            if dist[i, j] < 0:
                dist[i, j] += 100
    return dist

@nb.njit(parallel=False)
def func_jit3(a, b):
    dist = np.empty_like(a)
    for i in nb.prange(a.shape[0]):
        dist[i] = (a[i] - b[i]) + (a[i] - b[i] < 0)*100
    return dist

np.allclose(func_jit1(a, b), func_jit2(a, b))
np.allclose(func_jit1(a, b), func_jit3(a, b))

%timeit func_jit1(a, b)
%timeit func_jit2(a, b)
%timeit func_jit3(a, b)