Search code examples
pythonnumpynumba

Using numpy random choice in numba


I'm trying to rewrite some of my code using numba (version 0.52, on windows 10, 64-bit), however I get an error that I don't understand using numpy random choice. Numpy random choice should be working with numba if I don't use the probability option. Here is the code:

import numpy as np
from numba import jit

@jit(nopython=True)

def Calc():
    a = np.array([1, -1])
    size = [3, 3, 3]
    values = np.random.choice(a, size=size)

Calc()

I get the following error:

Traceback (most recent call last):
  File "\\uni.au.dk\Users\au684834\Documents\Python\Aarhus\Simulations\Ising\Tests\Numbaerrors.py", line 10, in <module>
    Calc()
  File "C:\Users\au684834\Miniconda3\lib\site-packages\numba\core\dispatcher.py", line 414, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\au684834\Miniconda3\lib\site-packages\numba\core\dispatcher.py", line 357, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
[1m[1m[1mNo implementation of function Function(<built-in method choice of numpy.random.mtrand.RandomState object at 0x000002B05203B340>) found for signature:
 
 >>> choice(array(int64, 1d, C), size=list(int64)<iv=[3, 3, 3]>)
 
There are 2 candidate implementations:
[1m  - Of which 1 did not match due to:
  Overload in function 'choice': File: numba\cpython\randomimpl.py: Line 1346.
    With argument(s): '(array(int64, 1d, C), size=list(int64)<iv=None>)':[0m
[1m   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   [1m[1m[1mNo implementation of function Function(<built-in function empty>) found for signature:
    
    >>> empty(list(int64)<iv=None>, class(int64))
    
   There are 2 candidate implementations:
   [1m      - Of which 2 did not match due to:
         Overload of function 'empty': File: numba\core\typing\npydecl.py: Line 507.
           With argument(s): '(list(int64)<iv=None>, class(int64))':[0m
   [1m       No match.[0m
   [0m
   [0m[1mDuring: resolving callee type: Function(<built-in function empty>)[0m
   [0m[1mDuring: typing of call at C:\Users\au684834\Miniconda3\lib\site-packages\numba\cpython\randomimpl.py (1403)
   [0m
   [1m
   File "..\Users\au684834\Miniconda3\lib\site-packages\numba\cpython\randomimpl.py", line 1403:[0m
   [1m        def choice_impl(a, size=None, replace=True):
               <source elided>
               if replace:
   [1m                out = np.empty(size, dtype)
   [0m                [1m^[0m[0m
[0m
  raised from C:\Users\au684834\Miniconda3\lib\site-packages\numba\core\typeinfer.py:1071
[1m  - Of which 1 did not match due to:
  Overload in function 'choice': File: numba\cpython\randomimpl.py: Line 1346.
    With argument(s): '(array(int64, 1d, C), size=list(int64)<iv=[3, 3, 3]>)':[0m
[1m   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   [1m[1m[1mNo implementation of function Function(<built-in function empty>) found for signature:
    
    >>> empty(list(int64)<iv=[3, 3, 3]>, class(int64))
    
   There are 2 candidate implementations:
   [1m      - Of which 2 did not match due to:
         Overload of function 'empty': File: numba\core\typing\npydecl.py: Line 507.
           With argument(s): '(list(int64)<iv=None>, class(int64))':[0m
   [1m       No match.[0m
   [0m
   [0m[1mDuring: resolving callee type: Function(<built-in function empty>)[0m
   [0m[1mDuring: typing of call at C:\Users\au684834\Miniconda3\lib\site-packages\numba\cpython\randomimpl.py (1403)
   [0m
   [1m
   File "..\Users\au684834\Miniconda3\lib\site-packages\numba\cpython\randomimpl.py", line 1403:[0m
   [1m        def choice_impl(a, size=None, replace=True):
               <source elided>
               if replace:
   [1m                out = np.empty(size, dtype)
   [0m                [1m^[0m[0m
[0m
  raised from C:\Users\au684834\Miniconda3\lib\site-packages\numba\core\typeinfer.py:1071
[0m
[0m[1mDuring: resolving callee type: Function(<built-in method choice of numpy.random.mtrand.RandomState object at 0x000002B05203B340>)[0m
[0m[1mDuring: typing of call at \\uni.au.dk\Users\au684834\Documents\Python\Aarhus\Simulations\Ising\Tests\Numbaerrors.py (8)
[0m
[1m
File "\\uni.au.dk\Users\au684834\Documents\Python\Aarhus\Simulations\Ising\Tests\Numbaerrors.py", line 8:[0m
[1mdef Calc():
    <source elided>
    size = [3, 3, 3]
[1m    values = np.random.choice(a, size=size)
[0m    [1m^[0m[0m

Not sure what I am doing wrong.


Solution

  • Use a tuple as size, not a list:

    @njit
    def calc():
        a = np.array([1, -1])
        size = (3, 3, 3)
        values = np.random.choice(a, size=size)
        return values
    

    Now:

    >>> calc()  # doctest: +SKIP
    array([[[-1,  1,  1],
            [-1, -1,  1],
            [-1, -1, -1]],
    
           [[-1, -1,  1],
            [-1, -1, -1],
            [-1, -1, -1]],
    
           [[ 1,  1,  1],
            [ 1,  1,  1],
            [ 1, -1,  1]]])
    

    Note however that not all numpy functions are supported by numba. One example among many is np.clip().