inconsistent array order for bare and qualified jit decorator

[yt87]

Consider the following code:

import functools
import numpy as np
import numba as nb

N = 4
M = 2
a = np.arange(M*N, dtype=np.float64).reshape(M, N).T
b = np.ones(M, dtype=np.float64)

def calc0(x, y, out=None):
    op = np.multiply
    it = np.nditer((x, y, None), [], [['readonly'], ['readonly'], ['writeonly', 'allocate']])
    for (a, b, c) in it:                                
        op(a, b, out=c)
    return it.operands[2]

@nb.njit
def calc1(x, y, out):
    it = np.nditer((x, y))
    for i, (u, v) in enumerate(it):
        out[i] = u * v
        
@nb.njit(nb.void(nb.float64[:,:], nb.float64[:], nb.float64[:]))
def calc2(x, y, out):
    it = np.nditer((x, y))
    for i, (u, v) in enumerate(it):
        out[i] = u * v
        
def calc1_wrap(x, y):
    shape = np.broadcast_shapes(x.shape, y.shape)
    size = functools.reduce(lambda u, v: u*v, shape)
    out = np.empty(size, dtype=np.float64)
    calc1(x, y, out)
    return out.reshape(shape)
        
def calc2_wrap(x, y):
    shape = np.broadcast_shapes(x.shape, y.shape)
    size = functools.reduce(lambda u, v: u*v, shape)
    out = np.empty(size, dtype=np.float64)
    calc2(x, y, out)
    return out.reshape(shape)

a * b
calc0(a, b)
calc1_wrap(a, b)
calc2_wrap(a, b)
        
calc1.signatures
calc2.signatures

The output is:

array([[0., 4.],
       [1., 5.],
       [2., 6.],
       [3., 7.]])

array([[0., 4.],
       [1., 5.],
       [2., 6.],
       [3., 7.]])

array([[0., 1.],
       [2., 3.],
       [4., 5.],
       [6., 7.]])

array([[0., 4.],
       [1., 5.],
       [2., 6.],
       [3., 7.]])

[(array(float64, 2d, F), array(float64, 1d, C), array(float64, 1d, C))]
[(array(float64, 2d, A), array(float64, 1d, A), array(float64, 1d, A))]

The third array stands out. It seems that the bare @njit assumes that the a is in Fortran order. This is not consistent with numpy behaviour. Is there a reason for this, or a minor bug?

One could argue that numpy documentation is misleading: https://numpy.org/doc/stable/reference/generated/numpy.nditer.html?highlight=nditer#numpy.nditer

[stuartarchibald]

Thanks for the report. I think the issue is in the 4th array arising from the call to calc2_wrap. Running the same script with JIT compilation disabled (via setting the environment variable NUMBA_DISABLE_JIT=1) yields:

[[0. 4.]
 [1. 5.]
 [2. 6.]
 [3. 7.]]
[[0. 4.]
 [1. 5.]
 [2. 6.]
 [3. 7.]]
[[0. 1.]
 [2. 3.]
 [4. 5.]
 [6. 7.]]
[[0. 1.]
 [2. 3.]
 [4. 5.]
 [6. 7.]]

whereas with the JIT enabled it produces:

[[0. 4.]
 [1. 5.]
 [2. 6.]
 [3. 7.]]
[[0. 4.]
 [1. 5.]
 [2. 6.]
 [3. 7.]]
[[0. 1.]
 [2. 3.]
 [4. 5.]
 [6. 7.]]
[[0. 4.]
 [1. 5.]
 [2. 6.]
 [3. 7.]]

I suspect the nditer isn't matching NumPy in the case of an 'A' array order, it would require a runtime check to determine the order ahead of iteration.

MWR for the issue:

from numba import njit, float64
import numpy as np

N = 4
M = 2
a = np.arange(M * N, dtype=np.float64).reshape(M, N).T

@njit((float64[:,:],))
def foo(x):
    for i in np.nditer(x):
        print(i)

print('Python:')
foo.py_func(a)
print('\nJIT:')
foo(a)

This produces:

Python:
0.0
1.0
2.0
3.0
4.0
5.0
6.0
7.0

JIT:
0.0
4.0
1.0
5.0
2.0
6.0
3.0
7.0