slicing c-contiguous array in row-major way yields non-c-contiguous array in numba (though not in numpy)

I have a four-d array (let’s call it four_d_arr) that is C-contiguous. In Python, slicing the array like three_d_arr = four_d_arr[0,:,:,:] yields a three dimensional array that’s C-contiguous. However, this behavior is not the same in numba (and I’m not sure why?). I detected this behavior when trying to reshape the three-d array from the slice (reshape() supports contiguous array only). I included code below that shows how such a slice maintains C-contiguity in numpy, but when compiled with numba the slice is no longer contiguous. I am able to reshape the three_d array by copying it with np.ascontiguousarray, which seems to indicate that the slicing process in numba is breaking the c-contiguity…

Here’s a minimal script that reproduces the problem:

import numpy as np                                                                                                                                                                                                                                                                                                           
import numba                                                                                                                                                                                                                                                                                                                 
from numba import jit                                                                                                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                                                                             
def reshape_arr(four_d_array):                                                                                                                                                                                                                                                                                               
    """                                                                                                                                                                                                                                                                                                                      
    Demonstrate that a sliced four d array along its first                                                                                                                                                                                                                                                                   
    dimension is still C contiguous                                                                                                                                                                                                                                                                                          
    """                                                                                                                                                                                                                                                                                                                      
    print('Input array flags\n', four_d_array.flags)                                                                                                                                                                                                                                                                         
    three_d_array = four_d_array[0,:,:,:]                                                                                                                                                                                                                                                                                    
    print('Sliced array flags (sliced on the first axis, so should still be c contiguous)\n')                                                                                                                                                                                                                                
    print(three_d_array.flags)                                                                                                                                                                                                                                                                                               
    last_axis_dim = three_d_array.shape[1]*three_d_array.shape[2]                                                                                                                                                                                                                                                            
    first_axis_dim = three_d_array.shape[0]                                                                                                                                                                                                                                                                                  
    three_d_array = three_d_array.reshape(first_axis_dim, last_axis_dim)                                                                                                                                                                                                                                                     
    return                                                                                                                                                                                                                                                                                                                   
                                                                                                                                                                                                                                                                                                                             
@jit(nopython=True)                                                                                                                                                                                                                                                                                                          
def jit_reshape_arr_as_cont(four_d_array):                                                                                                                                                                                                                                                                                   
    """                                                                                                                                                                                                                                                                                                                      
    Demonstrate that copying the array with                                                                                                                                                                                                                                                                                  
    np.ascontiguousarray permits successful                                                                                                                                                                                                                                                                                  
    compilation                                                                                                                                                                                                                                                                                                              
    """                                                                                                                                                                                                                                                                                                                      
    three_d_array = four_d_array[0,:,:,:]                                                                                                                                                                                                                                                                                    
    last_axis_dim = three_d_array.shape[1]*three_d_array.shape[2]                                                                                                                                                                                                                                                            
    first_axis_dim = three_d_array.shape[0]                                                                                                                                                                                                                                                                                  
    three_d_array = np.ascontiguousarray(three_d_array)                                                                                                                                                                                                                                                                      
    three_d_array = three_d_array.reshape(first_axis_dim, last_axis_dim)                                                                                                                                                                                                                                                     
    return                                                                                                                                                                                                                                                                                                                   
                                                                                                                                                                                                                                                                                                                             
@jit(nopython=True)                                                                                                                                                                                                                                                                                                          
def jit_reshape_arr(four_d_array):                                                                                                                                                                                                                                                                                           
    """                                                                                                                                                                                                                                                                                                                      
    Demonstrate that the sliced array fails                                                                                                                                                                                                                                                                                  
    to be reshaped, apparently because the                                                                                                                                                                                                                                                                                   
    C-contiguity is broken by slicing                                                                                                                                                                                                                                                                                        
    """                                                                                                                                                                                                                                                                                                                      
    three_d_array = four_d_array[0,:,:,:]                                                                                                                                                                                                                                                                                    
    last_axis_dim = three_d_array.shape[1]*three_d_array.shape[2]                                                                                                                                                                                                                                                            
    first_axis_dim = three_d_array.shape[0]                                                                                                                                                                                                                                                                                  
    three_d_array = three_d_array.reshape(first_axis_dim, last_axis_dim)                                                                                                                                                                                                                                                     
    return                                                                                                                                                                                                                                                                                                                   
                                                                                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                                                                             
if __name__ == '__main__':                                                                                                                                                                                                                                                                                                   
    arr = np.ones((3,4,5,6))                                                                                                                                                                                                                                                                                                 
    print(arr.flags)                                                                                                                                                                                                                                                                                                         
    print('Run without jit. Note that at all stages of this function the array is c contiguous')                                                                                                                                                                                                                             
    reshape_arr(arr)                                                                                                                                                                                                                                                                                                         
    print('Run with jit, but copy sliced array with "as contiguous array". It succeeds')                                                                                                                                                                                                                                     
    jit_reshape_arr_as_cont(arr)                                                                                                                                                                                                                                                                                             
    print('Now run with jit, but without copying array. reshape throws error \'reshape() supports contiguous arrays only\'')                                                                                                                                                                                                                          
    jit_reshape_arr(arr)                    

---

The output from running this on my machine yields

>>> python bug_test.py                                                                                                                                                            
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

Run without jit. Note that at all stages of this function the array is c contiguous Input array flags

   C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

Sliced array flags (sliced on the first axis, so should still be c contiguous)

  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : False
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

Run with jit, but copy sliced array with “as contiguous array”. It succeeds Now run with jit, but without copying array. reshape throws error 'reshape() supports contiguous arrays only’

Traceback (most recent call last):
  File "bug_test.py", line 42, in <module>
    jit_reshape_arr(arr)
  File "/home/hunter/anaconda3/lib/python3.8/site-packages/numba/core/dispatcher.py", line 415, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/hunter/anaconda3/lib/python3.8/site-packages/numba/core/dispatcher.py", line 358, in error_rewrite
    reraise(type(e), e, None)
  File "/home/hunter/anaconda3/lib/python3.8/site-packages/numba/core/utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
reshape() supports contiguous array only
- Resolution failure for non-literal arguments:
None

During: resolving callee type: BoundFunction(array.reshape for array(float64, 3d, A))
During: typing of call at bug_test.py (29)


File "bug_test.py", line 29:
def jit_reshape_arr(four_d_array):
    <source elided>
    first_axis_dim = three_d_array.shape[0]
    three_d_array = three_d_array.reshape(first_axis_dim, last_axis_dim)

[hunterakins]

I believe the problem is that keeps_contiguity only permits slices in the single innermost index. This is the correct behavior except in the special case that the last N slices are requesting the full axis. However, this corresponds to the ellipsis notation. Therefore, the problem is resolved **if the user passes in ellipsis instead of the tuple of slices **([int, …] instead of [int, :,:,:]). I considered ellipsis to be an optional notational convenience but it seems it has a sort of fundamental role. Here’s a function that shows this works (you can plug this into the minimal working example above)

@jit(nopython=True)                                                                                                                                                                                         
def jit_reshape_arr_with_ell(four_d_array):                                                                                                                                                                 
    """                                                                                                                                                                                                     
    Demonstrate that using ellipsis works                                                                                                                                                                   
    """                                                                                                                                                                                                     
    three_d_array = four_d_array[0,...]     # index with ellipsis                                                                                                                                                                
    last_axis_dim = three_d_array.shape[1]*three_d_array.shape[2]                                                                                                                                           
    first_axis_dim = three_d_array.shape[0]                                                                                                                                                                 
    three_d_array = three_d_array.reshape(first_axis_dim, last_axis_dim)                                                                                                                                    
    return

I think a nice solution on the back end could be to cast trailing tuples of full-axis slices to an ellipsis.

[jesseholzerpnnl]

@hunterakins Thanks for reporting this and for figuring out a solution. I needed that solution and would never in a million years have found it otherwise.