0

I know that numba and not-fixed-dimensional arrays do not go particularly well together. In my case I know that the dimension of the array will be 1,2 or 3. Therefore I could always write every method 3 times like this:

import numba
import numpy as np

def pick1D(arr:np.ndarray):
    i=np.random.randint(0,arr.shape[0])
    return arr[i]

def pick2D(arr:np.ndarray):
    i=np.random.randint(0,arr.shape[0])
    j=np.random.randint(0,arr.shape[1])
    return arr[i,j]

def pick3D(arr:np.ndarray):
    i=np.random.randint(0,arr.shape[0])
    j=np.random.randint(0,arr.shape[1])
    k=np.random.randint(0,arr.shape[2])
    return arr[i,j,k]

@numba.generated_jit
def ndpick(arr:np.ndarray):
    if arr.ndim==1:
        return pick1D
    elif arr.ndim==2:
        return pick2D
    elif arr.ndim==3:
        return pick3D
    return 0 #whatever

array=np.arange(24).reshape((2,3,4))
ndpick(array)

In pure python I could do some index-magic like this:

def pickND(arr):
    index = tuple(np.random.randint(0,arr.shape))
    return arr[index]

but as Numba does not support the tuple-constructor this does not work.

As my (real) methods contain a lot more code which is exactly the same: Is there any way to prevent copy-pasting the code 3 times? I think of something like templates from C++.

meneken17
  • 350
  • 1
  • 10
  • 1
    Does this answer your question? [Index an array with an array in numba](https://stackoverflow.com/questions/71977824/index-an-array-with-an-array-in-numba) – Jérôme Richard Jun 19 '22 at 16:51
  • @JérômeRichard No. I know the problem. I am looking for an approach to avoid coding the (almost) same method 3 times. – meneken17 Jun 19 '22 at 17:13
  • 1
    This is what the answer does. There is no templates in Numba. Only functions and a basic typing system. Variables needs to be typed and, as specified in the link, array dimensions are part of the type. So you need 3 instantiation of the function (mandatory in Numba). AFAIK, the only way to index a ND Numba arrays in a generic way is to flatten it because of the typing system. The answer is sufficient to make a unique generic function the way you describe. Furthermore, the example of the OP is very similar to your question (1D/2D/3D arrays with the same note about using a tuple). – Jérôme Richard Jun 19 '22 at 18:35

0 Answers0