Numba TypingError - too many positional arguments when trying to load saved umap models with v0.5.13

[dmyung]

Version 0.5.13 with numba 0.60.0 and 0.59.1 present this error when trying to load saved umap pickles when running plain umap with Jaccard on scipy sparse vectors.

pynndescent==0.5.13
numba==0.59.1 & 0.6.0 I receive this error on loading a pickle:

UMAP(angular_rp_forest=True, metric='jaccard', min_dist=0.1, n_jobs=1, n_neighbors=25, random_state=42, verbose=True)

...

with open('pickled_umap.pkl', 'wb') as fout:
    pickle.dump(fit_model, fout)


with open('pickled_umap.pkl', 'rb') as fin:
    foo = pickle.load(fin)

Building and compiling sparse search function
*** numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Internal error at <numba.core.typeinfer.CallConstraint object at 0x76f31c5e55d0>.
too many positional arguments
During: resolving callee type: type(CPUDispatcher(<function alternative_jaccard at 0x76f3b287ca60>))
During: typing of call at ../lib/python3.10/site-packages/pynndescent/pynndescent_.py (1560)

Enable logging at debug level for details.

File "../../../../../../miniforge3/envs/chemspace_env/lib/python3.10/site-packages/pynndescent/pynndescent_.py", line 1560:
        def search_closure(
            <source elided>
                    d = np.float32(
                        dist(
                        ^

(0.59.x is more detailed, 0.60.0 errors were not as helfpul)

With 0.5.12 and 0.5.11 and 0.59.1 the save and load work as expected.

A few other notes and observations.

I went down to pynndescent=0.5.1 and at that point umap is unimportable presenting pynndescent->numba related bootstrapping issues:

  File "~/umap-develop/src/umap/step/cluster/runner_umap.py", line 15, in <module>
    import umap
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/umap/__init__.py", line 2, in <module>
    from .umap_ import UMAP
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/umap/umap_.py", line 48, in <module>
    from pynndescent import NNDescent
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/pynndescent/__init__.py", line 3, in <module>
    from .pynndescent_ import NNDescent, PyNNDescentTransformer
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/pynndescent/pynndescent_.py", line 21, in <module>
    import pynndescent.sparse as sparse
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/pynndescent/sparse.py", line 10, in <module>
    from pynndescent.utils import norm, tau_rand
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/pynndescent/utils.py", line 630, in <module>
    def simple_heap_push(priorities, indices, p, n):
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/numba/core/decorators.py", line 236, in wrapper
    disp.compile(sig)
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/numba/core/dispatcher.py", line 957, in compile
    cres = self._compiler.compile(args, return_type)
  File "/home/ubuntu/miniforge3/envs/umap_env/lib/python3.10/site-packages/numba/core/dispatcher.py", line 125, in compile
    status, retval = self._compile_cached(args, return_type)
...

AssertionError: Failed in nopython mode pipeline (step: native lowering)
Storing i64 to ptr of i16 ('size'). FE type uint16

Workaround

As of this writing, I am able to confirm and reliably pickle and joblib.dump()/load() with pynndescent=0.5.12 and as numba=0.60.0, the current installation recipes from conda for umap install 0.5.13 and 0.60. Setting a minimum numba compatibility with pynndescent may help as well as seen from the other import issues and even pickle output issues I ran into during my testing. Numba is mentioned here as the nature of the error messages changed as I downgraded versions, but upon writing this issue, it seems like going from 59 to 60 did nothing to break things, except some behavior in messaging was changed.

[lmcinnes]

I'll have to figure out what has gone astray here. It isn't that obvious immediately. Thanks for at least documenting the workaround.