numba fails to resolve run_quicksort in pynndescent/sparse.py [umap-learn]

Hey there,

I'm not sure if this is the right place to ask... I'm trying to get umap-learn to run on archlinux with python 3.9 and llvm10. Here are the dependencies from umap-learn that I installed via pip:

Requirement already satisfied: umap-learn in /home/milian/.local/lib/python3.9/site-packages (0.5.1)
Requirement already satisfied: numpy>=1.17 in /usr/lib/python3.9/site-packages (from umap-learn) (1.20.0)
Requirement already satisfied: scikit-learn>=0.22 in /home/milian/.local/lib/python3.9/site-packages (from umap-learn) (0.24.1)
Requirement already satisfied: scipy>=1.0 in /home/milian/.local/lib/python3.9/site-packages (from umap-learn) (1.6.0)
Requirement already satisfied: numba>=0.49 in /home/milian/.local/lib/python3.9/site-packages (from umap-learn) (0.51.2)
Requirement already satisfied: pynndescent>=0.5 in /home/milian/.local/lib/python3.9/site-packages (from umap-learn) (0.5.2)
Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /home/milian/.local/lib/python3.9/site-packages (from numba>=0.49->umap-learn) (0.34.0)
Requirement already satisfied: setuptools in /usr/lib/python3.9/site-packages (from numba>=0.49->umap-learn) (53.0.0)
Requirement already satisfied: joblib>=0.11 in /home/milian/.local/lib/python3.9/site-packages (from pynndescent>=0.5->umap-learn) (1.0.1)
Requirement already satisfied: threadpoolctl>=2.0.0 in /home/milian/.local/lib/python3.9/site-packages (from scikit-learn>=0.22->umap-learn) (2.1.0)

But when I run a umap computation, then I get this wall of text as an error message:

Traceback (most recent call last):
  File ".../_UMAPDataGenerator.py", line 59, in <module>
    import umap  # noqa
  File "/home/milian/.local/lib/python3.9/site-packages/umap/__init__.py", line 2, in <module>
    from .umap_ import UMAP
  File "/home/milian/.local/lib/python3.9/site-packages/umap/umap_.py", line 47, in <module>
    from pynndescent import NNDescent
  File "/home/milian/.local/lib/python3.9/site-packages/pynndescent/__init__.py", line 3, in <module>
    from .pynndescent_ import NNDescent, PyNNDescentTransformer
  File "/home/milian/.local/lib/python3.9/site-packages/pynndescent/pynndescent_.py", line 21, in <module>
    import pynndescent.sparse as sparse
  File "/home/milian/.local/lib/python3.9/site-packages/pynndescent/sparse.py", line 343, in <module>
    def sparse_alternative_jaccard(ind1, data1, ind2, data2):
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/decorators.py", line 218, in wrapper
    disp.compile(sig)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/dispatcher.py", line 819, in compile
    cres = self._compiler.compile(args, return_type)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/dispatcher.py", line 82, in compile
    raise retval
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/dispatcher.py", line 92, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/dispatcher.py", line 105, in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler.py", line 627, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler.py", line 363, in compile_extra
    return self._compile_bytecode()
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler.py", line 425, in _compile_bytecode
    return self._compile_core()
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler.py", line 405, in _compile_core
    raise e
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler.py", line 396, in _compile_core
    pm.run(self.state)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 341, in run
    raise patched_exception
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 332, in run
    self._runPass(idx, pass_inst, state)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 291, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 264, in check
    mangled = func(compiler_state)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/typed_passes.py", line 92, in run_pass
    typemap, return_type, calltypes = type_inference_stage(
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/typed_passes.py", line 70, in type_inference_stage
    infer.propagate(raise_errors=raise_errors)
  File "/home/milian/.local/lib/python3.9/site-packages/numba/core/typeinfer.py", line 1071, in propagate
    raise errors[0]
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function make_quicksort_impl.<locals>.run_quicksort at 0x7f6312ddf790>) found for signature:
 
 >>> run_quicksort(array(int32, 1d, C))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'register_jitable.<locals>.wrap.<locals>.ov_wrap': File: numba/core/extending.py: Line 150.
    With argument(s): '(array(int32, 1d, C))':
   Rejected as the implementation raised a specific error:
     UnsupportedError: Failed in nopython mode pipeline (step: analyzing bytecode)
   Use of unsupported opcode (LOAD_ASSERTION_ERROR) found
   
   File "../../../../.local/lib/python3.9/site-packages/numba/misc/quicksort.py", line 180:
       def run_quicksort(A):
           <source elided>
               while high - low >= SMALL_QUICKSORT:
                   assert n < MAX_STACK
                   ^
   
  raised from /home/milian/.local/lib/python3.9/site-packages/numba/core/byteflow.py:269

During: resolving callee type: Function(<function make_quicksort_impl.<locals>.run_quicksort at 0x7f6312ddf790>)
During: typing of call at /home/milian/.local/lib/python3.9/site-packages/numba/np/arrayobj.py (5007)


File "../../../../.local/lib/python3.9/site-packages/numba/np/arrayobj.py", line 5007:
    def array_sort_impl(arr):
        <source elided>
        # Note we clobber the return value
        sort_func(arr)
        ^

During: lowering "$14call_method.5 = call $12load_method.4(func=$12load_method.4, args=[], kws=(), vararg=None)" at /home/milian/.local/lib/python3.9/site-packages/numba/np/arrayobj.py (5017)
During: lowering "$8call_method.3 = call $4load_method.1(arr, func=$4load_method.1, args=[Var(arr, sparse.py:28)], kws=(), vararg=None)" at /home/milian/.local/lib/python3.9/site-packages/pynndescent/sparse.py (28)
During: resolving callee type: type(CPUDispatcher(<function arr_unique at 0x7f6312e59040>))
During: typing of call at /home/milian/.local/lib/python3.9/site-packages/pynndescent/sparse.py (41)

During: resolving callee type: type(CPUDispatcher(<function arr_unique at 0x7f6312e59040>))
During: typing of call at /home/milian/.local/lib/python3.9/site-packages/pynndescent/sparse.py (41)


File "../../../../.local/lib/python3.9/site-packages/pynndescent/sparse.py", line 41:
def arr_union(ar1, ar2):
    <source elided>
    else:
        return arr_unique(np.concatenate((ar1, ar2)))
        ^

During: resolving callee type: type(CPUDispatcher(<function arr_union at 0x7f6312e59310>))
During: typing of call at /home/milian/.local/lib/python3.9/site-packages/pynndescent/sparse.py (344)

During: resolving callee type: type(CPUDispatcher(<function arr_union at 0x7f6312e59310>))
During: typing of call at /home/milian/.local/lib/python3.9/site-packages/pynndescent/sparse.py (344)


File "../../../../.local/lib/python3.9/site-packages/pynndescent/sparse.py", line 344:
def sparse_alternative_jaccard(ind1, data1, ind2, data2):
    num_non_zero = arr_union(ind1, ind2).shape[0]
    ^

Can anyone tell me where to hunt for the problem here? Is that an incompatibility with llvm? Is something outdated on pip? Anything else I could do to fix this?

Many thanks

Begin by reproducing the reported import or UMAP computation with Python 3.9 and the listed dependency versions. Read the traceback locations in pynndescent/sparse.py, numba/misc/quicksort.py, and numba/np/arrayobj.py to determine whether the failure is a dependency compatibility issue. Done means identifying a confirmed cause and a verified compatible version or repository change.