Indexing with array gives wrong result for negative entries

[jamesavery]

Example:

a  = np.arange(1,5);
ix = a[:,None]-a[None,:];
print(a[ix])

Numpy result:

[[1 4 3 2]
 [2 1 4 3]
 [3 2 1 4]
 [4 3 2 1]]

Bohrium result:

[[1 0 0 0]
 [2 1 0 0]
 [3 2 1 0]
 [4 3 2 1]]

[jamesavery]

The following user kernel is a simple implementation of correct indexing with arrays (any dimension over any axis):

ctype = {
    'float32': 'float',
    'float64': 'double',
    'complex64': 'float complex',
    'complex128':'double complex',
    'int8':  'int8_t',
    'int16': 'int16_t',
    'int32': 'int32_t',
    'int64': 'int64_t',
    'uint8':  'uint8_t',
    'uint16': 'uint16_t',
    'uint32': 'uint32_t',
    'uint64': 'uint64_t',    
    'bool': 'uint8_t'
}

   
def axis_split(A,axis):
    (Am,Ai,An) = (product(A.shape[:axis]),  A.shape[axis], product(A.shape[axis+1:]));
    return (Am,Ai,An)
    
def take_cpu(A,I,axis=0):        
    (Am,Ai,An) = axis_split(A,axis)
    Ar = bh.user_kernel.make_behaving(A.reshape(Am,Ai,An),dtype=A.dtype);
    I  = bh.user_kernel.make_behaving(I,dtype=np.int64);
    
    Ii = product(I.shape)
    
    new_shape = list(A.shape[:axis])+list(I.shape)+list(A.shape[axis+1:]);
    cmd = bh.user_kernel.get_default_compiler_command()
    kernel = read_file("lookup-cpu.c") % {'ctype':ctype[A.dtype.name],'Am':Am,'Ai':Ai,'Ii':Ii,'An':An};
       
    AI = bh.empty(new_shape,dtype=A.dtype)
    bh.user_kernel.execute(kernel, [AI, A, I], compiler_command=cmd)
    
    return AI.reshape(new_shape);

The ctype dict is generally useful and could be added to bh.user_kernel.

Contents of lookup-cpu.c:

#include <stdint.h>
#include <stdio.h>
#include <complex.h>

typedef %(ctype)s scalar_t;

/* A: Am x Ai x An -> Am x Ii x An */
void index_axis(scalar_t *O,
		scalar_t *A, uint64_t Am, uint64_t Ai, uint64_t An,
		int64_t *Index, uint64_t Ii)
{
#pragma omp parallel for collapse(3)
  for(uint64_t i=0;i<Am;i++)
    for(uint64_t j=0;j<Ii;j++)
      for(uint64_t k=0;k<An;k++){
	uint64_t I_j = Index[j]>=0? Index[j] : Ai+Index[j];
	O[k+j*An+i*Ii*An] = A[k+I_j*An+i*Ai*An];
      }
}



void execute(scalar_t *O, scalar_t *A, int64_t *Index)
{
  uint64_t Am = %(Am)d, Ai = %(Ai)d, Ii = %(Ii)d, An =%(An)d;
  return index_axis(O, A,Am,Ai,An, Index,Ii);
}