Bug fitting umap with numba

[timonmerk]

I see an error when importing and fitting joblib (when executing import from vscode debug console):

joblib.externals.loky.process_executor._RemoteTraceback:
"""
Traceback (most recent call last):
File "_pydevd_sys_monitoring\_pydevd_sys_monitoring_cython.pyx", line 592, in _pydevd_sys_monitoring_cython._get_func_code_info
KeyError: (1408, '/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/rp_trees.py', <code object get_leaves_from_tree at 0x3b1d21210, file "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/rp_trees.py", line 1408>)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/_utils.py", line 72, in call
return self.func(**kwargs)
^^^^^^^^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 598, in call
return [func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "", line 69, in cfunc.to_py.__Pyx_CFunc_893235__29_pydevd_sys_monitoring_cython_object__lParen__etc_to_py_4code_18instruction_offset.wrap
File "_pydevd_sys_monitoring\_pydevd_sys_monitoring_cython.pyx", line 1697, in _pydevd_sys_monitoring_cython._start_method_event
File "_pydevd_sys_monitoring\_pydevd_sys_monitoring_cython.pyx", line 599, in _pydevd_sys_monitoring_cython._get_func_code_info
AssertionError: <code object call at 0x12e1e1930, file "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 594> != <code object get_leaves_from_tree at 0x3b1d21210, file "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/rp_trees.py", line 1408>
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "", line 1, in
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/umap/umap_.py", line 2928, in fit_transform
self.fit(X, y, force_all_finite, **kwargs)
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/umap/umap_.py", line 2635, in fit
) = nearest_neighbors(
^^^^^^^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/umap/umap_.py", line 330, in nearest_neighbors
knn_search_index = NNDescent(
^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/pynndescent_.py", line 806, in init
leaf_array = rptree_leaf_array(self._rp_forest)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/rp_trees.py", line 1436, in rptree_leaf_array
return np.vstack(rptree_leaf_array_parallel(rp_forest))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/rp_trees.py", line 1428, in rptree_leaf_array_parallel
result = joblib.Parallel(n_jobs=-1, require="sharedmem")(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 2007, in call
return output if self.return_generator else list(output)
^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 1650, in _get_outputs
yield from self._retrieve()
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 1754, in _retrieve
self._raise_error_fast()
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 1789, in _raise_error_fast
error_job.get_result(self.timeout)
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 745, in get_result
return self._return_or_raise()
^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 763, in _return_or_raise
raise self._result
AssertionError: <code object call at 0x12e1e1930, file "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/joblib/parallel.py", line 594> != <code object get_leaves_from_tree at 0x3b1d21210, file "/Users/Timon/Documents/mvts_transformer/.venv/lib/python3.12/site-packages/pynndescent/rp_trees.py", line 1408>

I am pretty sure that this is an incompatibility issue within numba, since the line in rp_trees.py uses numba: https://github.com/lmcinnes/pynndescent/blob/17e6c7046cdea5747aa0b7db21f66eb183f6c753/pynndescent/rp_trees.py#L1408

I tried to install different versions of numba and pynndescent, but couldn't fix the issue.

Python version: 3.12.6 on MacM3
umap version: 0.5.7
pynndescent version: 0.5.13
numba version: 0.60.0
joblib version: 1.4.2