Compute Transition Matrix Error CellRank Meets RNA Velocity

[anthonyinsalaco1]

When I try to run the code below on the tutorial, I get this error.

``Computing transition matrix using'deterministic'` model 0% 0/2531 [00:00<?, ?cell/s]

KeyError Traceback (most recent call last) File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:469, in IndexDataCacheFile.save(self, key, data) 467 try: 468 # If key already exists, we will overwrite the file --> 469 data_name = overloads[key] 470 except KeyError: 471 # Find an available name for the data file

KeyError: ((Function(<function norm at 0x7f7d371c2af0>), Literalint, Array(float64, 2, 'A', False, aligned=True)), ('x86_64-apple-darwin23.4.0', 'skylake', '+64bit,+adx,+aes,-amx-bf16,-amx-int8,-amx-tile,+avx,+avx2,-avx512bf16,-avx512bitalg,-avx512bw,-avx512cd,-avx512dq,-avx512er,-avx512f,-avx512fp16,-avx512ifma,-avx512pf,-avx512vbmi,-avx512vbmi2,-avx512vl,-avx512vnni,-avx512vp2intersect,-avx512vpopcntdq,-avxvnni,+bmi,+bmi2,-cldemote,+clflushopt,-clwb,-clzero,+cmov,+crc32,+cx16,+cx8,-enqcmd,+f16c,+fma,-fma4,+fsgsbase,+fxsr,-gfni,-hreset,+invpcid,-kl,-lwp,+lzcnt,+mmx,+movbe,-movdir64b,-movdiri,-mwaitx,+pclmul,-pconfig,-pku,+popcnt,-prefetchwt1,+prfchw,-ptwrite,-rdpid,+rdrnd,+rdseed,-rtm,+sahf,-serialize,+sgx,-sha,-shstk,+sse,+sse2,+sse3,+sse4.1,+sse4.2,-sse4a,+ssse3,-tbm,-tsxldtrk,-uintr,-vaes,-vpclmulqdq,-waitpkg,-wbnoinvd,-widekl,-xop,+xsave,+xsavec,+xsaveopt,+xsaves'), ('04d17047dec8a5ada0d1f7224d7e8baf78f8fb78c8913cbdf8a204a5c3f1492f', 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'))

During handling of the above exception, another exception occurred:

TypeError Traceback (most recent call last) Input In [23], in <cell line: 1>() ----> 1 vk.compute_transition_matrix()

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/_velocity_kernel.py:160, in VelocityKernel.compute_transition_matrix(self, model, backward_mode, similarity, softmax_scale, n_samples, seed, **kwargs) 157 return self 159 if softmax_scale is None: --> 160 softmax_scale = self._estimate_softmax_scale(backward_mode=backward_mode, similarity=similarity) 161 logg.info(f"Using softmax_scale={softmax_scale:.4f}") 162 params["softmax_scale"] = softmax_scale

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/_velocity_kernel.py:254, in VelocityKernel._estimate_softmax_scale(self, n_jobs, backend, kwargs) 247 def _estimate_softmax_scale( 248 self, 249 n_jobs: Optional[int] = None, 250 backend: Backend_t = DEFAULT_BACKEND, 251 kwargs, 252 ) -> float: 253 model = self._create_model(VelocityModel.DETERMINISTIC, softmaxscale=1.0, **kwargs) --> 254 , logits = model(n_jobs, backend) 255 return 1.0 / np.median(np.abs(logits.data))

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_velocity_model.py:64, in ModelABC.call(self, n_jobs, backend, show_progress_bar, kwargs) 56 def call( 57 self, 58 n_jobs: Optional[int] = None, (...) 61 kwargs: Any, 62 ) -> Tuple[np.ndarray, np.ndarray]: 63 ixs = self._ixs ---> 64 return parallelize( 65 self._compute_helper, 66 ixs, 67 n_jobs=n_jobs, 68 backend=backend, 69 show_progress_bar=show_progress_bar, 70 as_array=False, 71 extractor=lambda data: self._reconstruct_output(np.concatenate(data, axis=-1), ixs), 72 unit=self._unit, 73 )(**kwargs)

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/_utils/_parallelize.py:96, in parallelize..wrapper(*args, kwargs) 93 else: 94 pbar, queue, thread = None, None, None ---> 96 res = jl.Parallel(n_jobs=n_jobs, backend=backend)( 97 jl.delayed(callback)( 98 ((i, cs) if use_ixs else (cs,)), 99 args, 100 kwargs, 101 queue=queue, 102 ) 103 for i, cs in enumerate(collections) 104 ) 106 res = np.array(res) if as_array else res 107 if thread is not None:

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:1043, in Parallel.call(self, iterable) 1034 try: 1035 # Only set self._iterating to True if at least a batch 1036 # was dispatched. In particular this covers the edge (...) 1040 # was very quick and its callback already dispatched all the 1041 # remaining jobs. 1042 self._iterating = False -> 1043 if self.dispatch_one_batch(iterator): 1044 self._iterating = self._original_iterator is not None 1046 while self.dispatch_one_batch(iterator):

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:861, in Parallel.dispatch_one_batch(self, iterator) 859 return False 860 else: --> 861 self._dispatch(tasks) 862 return True

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:779, in Parallel._dispatch(self, batch) 777 with self._lock: 778 job_idx = len(self._jobs) --> 779 job = self._backend.apply_async(batch, callback=cb) 780 # A job can complete so quickly than its callback is 781 # called before we get here, causing self._jobs to 782 # grow. To ensure correct results ordering, .insert is 783 # used (rather than .append) in the following line 784 self._jobs.insert(job_idx, job)

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/_parallel_backends.py:208, in SequentialBackend.apply_async(self, func, callback) 206 def apply_async(self, func, callback=None): 207 """Schedule a func to be run""" --> 208 result = ImmediateResult(func) 209 if callback: 210 callback(result)

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/_parallel_backends.py:572, in ImmediateResult.init(self, batch) 569 def init(self, batch): 570 # Don't delay the application, to avoid keeping the input 571 # arguments in memory --> 572 self.results = batch()

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:262, in BatchedCalls.call(self) 258 def call(self): 259 # Set the default nested backend to self._backend but do not set the 260 # change the default number of processes to -1 261 with parallel_backend(self._backend, n_jobs=self._n_jobs): --> 262 return [func(*args, **kwargs) 263 for func, args, kwargs in self.items]

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:262, in (.0) 258 def call(self): 259 # Set the default nested backend to self._backend but do not set the 260 # change the default number of processes to -1 261 with parallel_backend(self._backend, n_jobs=self._n_jobs): --> 262 return [func(*args, **kwargs) 263 for func, args, kwargs in self.items]

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_velocity_model.py:85, in ModelABC._compute_helper(self, ixs, queue, kwargs) 82 neigh_ixs = indices[start:end] 83 n_neigh = len(neigh_ixs) ---> 85 ps, ls = self._compute(ix, neigh_ixs, kwargs) 86 if np.shape(ps) != (n_neigh,): 87 raise ValueError(f"Expected row of shape {(2, n_neigh)}, found {np.shape(ps)}.")

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_velocity_model.py:181, in Deterministic._compute(self, ix, neigh_ixs) 178 if self._backward_mode == BackwardMode.NEGATE: 179 v *= -1.0 --> 181 return self._similarity(v[None, :], W, self._softmax_scale)

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_similarity.py:238, in SimilarityHessian.call(self, v, D, softmax_scale) 226 @d.dedent 227 def call(self, v: np.ndarray, D: np.ndarray, softmax_scale: float = 1.0) -> Tuple[np.ndarray, np.ndarray]: 228 """%(sim_scheme.full_desc)s 229 230 Parameters (...) 236 %(sim_scheme.returns)s 237 """ # noqa: D400 --> 238 return _predict_transition_probabilities_numpy(v, D, softmax_scale, self._center_mean, self._scale_by_norm)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:442, in _DispatcherBase._compile_for_args(self, *args, **kws) 440 e.patch_message('\n'.join((str(e).rstrip(), help_msg))) 441 # ignore the FULL_TRACEBACKS config, this needs reporting! --> 442 raise e 443 finally: 444 self._types_active_call = []

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:375, in _DispatcherBase._compile_for_args(self, *args, **kws) 373 return_val = None 374 try: --> 375 return_val = self.compile(tuple(argtypes)) 376 except errors.ForceLiteralArg as e: 377 # Received request for compiler re-entry with the list of arguments 378 # indicated by e.requested_args. 379 # First, check if any of these args are already Literal-ized 380 already_lit_pos = [i for i in e.requested_args 381 if isinstance(args[i], types.Literal)]

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:905, in Dispatcher.compile(self, sig) 903 with ev.trigger_event("numba:compile", data=ev_details): 904 try: --> 905 cres = self._compiler.compile(args, return_type) 906 except errors.ForceLiteralArg as e: 907 def folded(args, kws):

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:80, in _FunctionCompiler.compile(self, args, return_type) 79 def compile(self, args, return_type): ---> 80 status, retval = self._compile_cached(args, return_type) 81 if status: 82 return retval

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:94, in _FunctionCompiler._compile_cached(self, args, return_type) 91 pass 93 try: ---> 94 retval = self._compile_core(args, return_type) 95 except errors.TypingError as e: 96 self._failed_cache[key] = e

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:107, in _FunctionCompiler._compile_core(self, args, return_type) 104 flags = self._customize_flags(flags) 106 impl = self._get_implementation(args, {}) --> 107 cres = compiler.compile_extra(self.targetdescr.typing_context, 108 self.targetdescr.target_context, 109 impl, 110 args=args, return_type=return_type, 111 flags=flags, locals=self.locals, 112 pipeline_class=self.pipeline_class) 113 # Check typing error if object mode is used 114 if cres.typing_error is not None and not flags.enable_pyobject:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:744, in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class) 720 """Compiler entry point 721 722 Parameter (...) 740 compiler pipeline 741 """ 742 pipeline = pipeline_class(typingctx, targetctx, library, 743 args, return_type, flags, locals) --> 744 return pipeline.compile_extra(func)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:438, in CompilerBase.compile_extra(self, func) 436 self.state.lifted = () 437 self.state.lifted_from = None --> 438 return self._compile_bytecode()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:506, in CompilerBase._compile_bytecode(self) 502 """ 503 Populate and run pipeline for bytecode input 504 """ 505 assert self.state.func_ir is None --> 506 return self._compile_core()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:481, in CompilerBase._compile_core(self) 478 except Exception as e: 479 if (utils.use_new_style_errors() and not 480 isinstance(e, errors.NumbaError)): --> 481 raise e 483 self.state.status.fail_reason = e 484 if is_final_pipeline:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:472, in CompilerBase._compile_core(self) 470 res = None 471 try: --> 472 pm.run(self.state) 473 if self.state.cr is not None: 474 break

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:364, in PassManager.run(self, state) 361 except Exception as e: 362 if (utils.use_new_style_errors() and not 363 isinstance(e, errors.NumbaError)): --> 364 raise e 365 msg = "Failed in %s mode pipeline (step: %s)" % \ 366 (self.pipeline_name, pass_desc) 367 patched_exception = self._patch_error(msg, e)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:356, in PassManager.run(self, state) 354 pass_inst = _pass_registry.get(pss).pass_inst 355 if isinstance(pass_inst, CompilerPass): --> 356 self._runPass(idx, pass_inst, state) 357 else: 358 raise BaseException("Legacy pass in use")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_lock.py:35, in _CompilerLock.call.._acquire_compile_lock(*args, kwargs) 32 @functools.wraps(func) 33 def _acquire_compile_lock(*args, *kwargs): 34 with self: ---> 35 return func(args, kwargs)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:311, in PassManager._runPass(self, index, pss, internal_state) 309 mutated |= check(pss.run_initialization, internal_state) 310 with SimpleTimer() as pass_time: --> 311 mutated |= check(pss.run_pass, internal_state) 312 with SimpleTimer() as finalize_time: 313 mutated |= check(pss.run_finalizer, internal_state)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:273, in PassManager._runPass..check(func, compiler_state) 272 def check(func, compiler_state): --> 273 mangled = func(compiler_state) 274 if mangled not in (True, False): 275 msg = ("CompilerPass implementations should return True/False. " 276 "CompilerPass with name '%s' did not.")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typed_passes.py:112, in BaseTypeInference.run_pass(self, state) 106 """ 107 Type inference and legalization 108 """ 109 with fallback_context(state, 'Function "%s" failed type inference' 110 % (state.func_id.func_name,)): 111 # Type inference --> 112 typemap, return_type, calltypes, errs = type_inference_stage( 113 state.typingctx, 114 state.targetctx, 115 state.func_ir, 116 state.args, 117 state.return_type, 118 state.locals, 119 raise_errors=self._raise_errors) 120 state.typemap = typemap 121 # save errors in case of partial typing

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typed_passes.py:93, in type_inference_stage(typingctx, targetctx, interp, args, return_type, locals, raise_errors) 91 infer.build_constraint() 92 # return errors in case of partial typing ---> 93 errs = infer.propagate(raise_errors=raise_errors) 94 typemap, restype, calltypes = infer.unify(raise_errors=raise_errors) 96 return _TypingResults(typemap, restype, calltypes, errs)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:1083, in TypeInferer.propagate(self, raise_errors) 1080 oldtoken = newtoken 1081 # Errors can appear when the type set is incomplete; only 1082 # raise them when there is no progress anymore. -> 1083 errors = self.constraints.propagate(self) 1084 newtoken = self.get_state_token() 1085 self.debug.propagate_finished()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:182, in ConstraintNetwork.propagate(self, typeinfer) 180 errors.append(utils.chain_exception(new_exc, e)) 181 elif utils.use_new_style_errors(): --> 182 raise e 183 else: 184 msg = ("Unknown CAPTURED_ERRORS style: " 185 f"'{config.CAPTURED_ERRORS}'.")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:160, in ConstraintNetwork.propagate(self, typeinfer) 157 with typeinfer.warnings.catch_warnings(filename=loc.filename, 158 lineno=loc.line): 159 try: --> 160 constraint(typeinfer) 161 except ForceLiteralArg as e: 162 errors.append(e)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:583, in CallConstraint.call(self, typeinfer) 581 fnty = typevars[self.func].getone() 582 with new_error_context("resolving callee type: {0}", fnty): --> 583 self.resolve(typeinfer, typevars, fnty)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:606, in CallConstraint.resolve(self, typeinfer, typevars, fnty) 604 fnty = fnty.instance_type 605 try: --> 606 sig = typeinfer.resolve_call(fnty, pos_args, kw_args) 607 except ForceLiteralArg as e: 608 # Adjust for bound methods 609 folding_args = ((fnty.this,) + tuple(self.args) 610 if isinstance(fnty, types.BoundFunction) 611 else self.args)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:1577, in TypeInferer.resolve_call(self, fnty, pos_args, kw_args) 1574 return sig 1575 else: 1576 # Normal non-recursive call -> 1577 return self.context.resolve_function_type(fnty, pos_args, kw_args)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typing/context.py:196, in BaseContext.resolve_function_type(self, func, args, kws) 194 # Prefer user definition first 195 try: --> 196 res = self._resolve_user_function_type(func, args, kws) 197 except errors.TypingError as e: 198 # Capture any typing error 199 last_exception = e

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typing/context.py:248, in BaseContext._resolve_user_function_type(self, func, args, kws, literals) 244 return self.resolve_function_type(func_type, args, kws) 246 if isinstance(func, types.Callable): 247 # XXX fold this into the call attribute logic? --> 248 return func.get_call_type(self, args, kws)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/types/functions.py:541, in Dispatcher.get_call_type(self, context, args, kws) 534 def get_call_type(self, context, args, kws): 535 """ 536 Resolve a call to this dispatcher using the given argument types. 537 A signature returned and it is ensured that a compiled specialization 538 is available for it. 539 """ 540 template, pysig, args, kws = \ --> 541 self.dispatcher.get_call_template(args, kws) 542 sig = template(context).apply(args, kws) 543 if sig:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:318, in _DispatcherBase.get_call_template(self, args, kws) 316 # Ensure an overload is available 317 if self._can_compile: --> 318 self.compile(tuple(args)) 320 # Create function type for typing 321 func_name = self.py_func.name

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:905, in Dispatcher.compile(self, sig) 903 with ev.trigger_event("numba:compile", data=ev_details): 904 try: --> 905 cres = self._compiler.compile(args, return_type) 906 except errors.ForceLiteralArg as e: 907 def folded(args, kws):

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:80, in _FunctionCompiler.compile(self, args, return_type) 79 def compile(self, args, return_type): ---> 80 status, retval = self._compile_cached(args, return_type) 81 if status: 82 return retval

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:94, in _FunctionCompiler._compile_cached(self, args, return_type) 91 pass 93 try: ---> 94 retval = self._compile_core(args, return_type) 95 except errors.TypingError as e: 96 self._failed_cache[key] = e

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:107, in _FunctionCompiler._compile_core(self, args, return_type) 104 flags = self._customize_flags(flags) 106 impl = self._get_implementation(args, {}) --> 107 cres = compiler.compile_extra(self.targetdescr.typing_context, 108 self.targetdescr.target_context, 109 impl, 110 args=args, return_type=return_type, 111 flags=flags, locals=self.locals, 112 pipeline_class=self.pipeline_class) 113 # Check typing error if object mode is used 114 if cres.typing_error is not None and not flags.enable_pyobject:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:744, in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class) 720 """Compiler entry point 721 722 Parameter (...) 740 compiler pipeline 741 """ 742 pipeline = pipeline_class(typingctx, targetctx, library, 743 args, return_type, flags, locals) --> 744 return pipeline.compile_extra(func)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:438, in CompilerBase.compile_extra(self, func) 436 self.state.lifted = () 437 self.state.lifted_from = None --> 438 return self._compile_bytecode()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:506, in CompilerBase._compile_bytecode(self) 502 """ 503 Populate and run pipeline for bytecode input 504 """ 505 assert self.state.func_ir is None --> 506 return self._compile_core()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:481, in CompilerBase._compile_core(self) 478 except Exception as e: 479 if (utils.use_new_style_errors() and not 480 isinstance(e, errors.NumbaError)): --> 481 raise e 483 self.state.status.fail_reason = e 484 if is_final_pipeline:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:472, in CompilerBase._compile_core(self) 470 res = None 471 try: --> 472 pm.run(self.state) 473 if self.state.cr is not None: 474 break

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:364, in PassManager.run(self, state) 361 except Exception as e: 362 if (utils.use_new_style_errors() and not 363 isinstance(e, errors.NumbaError)): --> 364 raise e 365 msg = "Failed in %s mode pipeline (step: %s)" % \ 366 (self.pipeline_name, pass_desc) 367 patched_exception = self._patch_error(msg, e)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:356, in PassManager.run(self, state) 354 pass_inst = _pass_registry.get(pss).pass_inst 355 if isinstance(pass_inst, CompilerPass): --> 356 self._runPass(idx, pass_inst, state) 357 else: 358 raise BaseException("Legacy pass in use")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_lock.py:35, in _CompilerLock.call.._acquire_compile_lock(*args, kwargs) 32 @functools.wraps(func) 33 def _acquire_compile_lock(*args, *kwargs): 34 with self: ---> 35 return func(args, kwargs)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:311, in PassManager._runPass(self, index, pss, internal_state) 309 mutated |= check(pss.run_initialization, internal_state) 310 with SimpleTimer() as pass_time: --> 311 mutated |= check(pss.run_pass, internal_state) 312 with SimpleTimer() as finalize_time: 313 mutated |= check(pss.run_finalizer, internal_state)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:273, in PassManager._runPass..check(func, compiler_state) 272 def check(func, compiler_state): --> 273 mangled = func(compiler_state) 274 if mangled not in (True, False): 275 msg = ("CompilerPass implementations should return True/False. " 276 "CompilerPass with name '%s' did not.")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typed_passes.py:112, in BaseTypeInference.run_pass(self, state) 106 """ 107 Type inference and legalization 108 """ 109 with fallback_context(state, 'Function "%s" failed type inference' 110 % (state.func_id.func_name,)): 111 # Type inference --> 112 typemap, return_type, calltypes, errs = type_inference_stage( 113 state.typingctx, 114 state.targetctx, 115 state.func_ir, 116 state.args, 117 state.return_type, 118 state.locals, 119 raise_errors=self._raise_errors) 120 state.typemap = typemap 121 # save errors in case of partial typing

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typed_passes.py:93, in type_inference_stage(typingctx, targetctx, interp, args, return_type, locals, raise_errors) 91 infer.build_constraint() 92 # return errors in case of partial typing ---> 93 errs = infer.propagate(raise_errors=raise_errors) 94 typemap, restype, calltypes = infer.unify(raise_errors=raise_errors) 96 return _TypingResults(typemap, restype, calltypes, errs)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:1083, in TypeInferer.propagate(self, raise_errors) 1080 oldtoken = newtoken 1081 # Errors can appear when the type set is incomplete; only 1082 # raise them when there is no progress anymore. -> 1083 errors = self.constraints.propagate(self) 1084 newtoken = self.get_state_token() 1085 self.debug.propagate_finished()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:182, in ConstraintNetwork.propagate(self, typeinfer) 180 errors.append(utils.chain_exception(new_exc, e)) 181 elif utils.use_new_style_errors(): --> 182 raise e 183 else: 184 msg = ("Unknown CAPTURED_ERRORS style: " 185 f"'{config.CAPTURED_ERRORS}'.")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:160, in ConstraintNetwork.propagate(self, typeinfer) 157 with typeinfer.warnings.catch_warnings(filename=loc.filename, 158 lineno=loc.line): 159 try: --> 160 constraint(typeinfer) 161 except ForceLiteralArg as e: 162 errors.append(e)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:583, in CallConstraint.call(self, typeinfer) 581 fnty = typevars[self.func].getone() 582 with new_error_context("resolving callee type: {0}", fnty): --> 583 self.resolve(typeinfer, typevars, fnty)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:606, in CallConstraint.resolve(self, typeinfer, typevars, fnty) 604 fnty = fnty.instance_type 605 try: --> 606 sig = typeinfer.resolve_call(fnty, pos_args, kw_args) 607 except ForceLiteralArg as e: 608 # Adjust for bound methods 609 folding_args = ((fnty.this,) + tuple(self.args) 610 if isinstance(fnty, types.BoundFunction) 611 else self.args)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:1577, in TypeInferer.resolve_call(self, fnty, pos_args, kw_args) 1574 return sig 1575 else: 1576 # Normal non-recursive call -> 1577 return self.context.resolve_function_type(fnty, pos_args, kw_args)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typing/context.py:196, in BaseContext.resolve_function_type(self, func, args, kws) 194 # Prefer user definition first 195 try: --> 196 res = self._resolve_user_function_type(func, args, kws) 197 except errors.TypingError as e: 198 # Capture any typing error 199 last_exception = e

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typing/context.py:248, in BaseContext._resolve_user_function_type(self, func, args, kws, literals) 244 return self.resolve_function_type(func_type, args, kws) 246 if isinstance(func, types.Callable): 247 # XXX fold this into the call attribute logic? --> 248 return func.get_call_type(self, args, kws)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/types/functions.py:541, in Dispatcher.get_call_type(self, context, args, kws) 534 def get_call_type(self, context, args, kws): 535 """ 536 Resolve a call to this dispatcher using the given argument types. 537 A signature returned and it is ensured that a compiled specialization 538 is available for it. 539 """ 540 template, pysig, args, kws = \ --> 541 self.dispatcher.get_call_template(args, kws) 542 sig = template(context).apply(args, kws) 543 if sig:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:318, in _DispatcherBase.get_call_template(self, args, kws) 316 # Ensure an overload is available 317 if self._can_compile: --> 318 self.compile(tuple(args)) 320 # Create function type for typing 321 func_name = self.py_func.name

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:912, in Dispatcher.compile(self, sig) 910 raise e.bind_fold_arguments(folded) 911 self.add_overload(cres) --> 912 self._cache.save_overload(sig, cres) 913 return cres.entry_point

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:652, in Cache.save_overload(self, sig, data) 648 """ 649 Save the data for the given signature in the cache. 650 """ 651 with self._guard_against_spurious_io_errors(): --> 652 self._save_overload(sig, data)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:662, in Cache._save_overload(self, sig, data) 660 key = self._index_key(sig, data.codegen) 661 data = self._impl.reduce(data) --> 662 self._cache_file.save(key, data)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:478, in IndexDataCacheFile.save(self, key, data) 476 break 477 overloads[key] = data_name --> 478 self._save_index(overloads) 479 self._save_data(data_name, data)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:522, in IndexDataCacheFile._save_index(self, overloads) 520 def _save_index(self, overloads): 521 data = self._source_stamp, overloads --> 522 data = self._dump(data) 523 with self._open_for_write(self._index_path) as f: 524 pickle.dump(self._version, f, protocol=-1)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:550, in IndexDataCacheFile._dump(self, obj) 549 def _dump(self, obj): --> 550 return dumps(obj)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/serialize.py:58, in dumps(obj) 56 with io.BytesIO() as buf: 57 p = pickler(buf, protocol=4) ---> 58 p.dump(obj) 59 pickled = buf.getvalue() 61 return pickled

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/cloudpickle/cloudpickle.py:1262, in Pickler.dump(self, obj) 1260 def dump(self, obj): 1261 try: -> 1262 return super().dump(obj) 1263 except RuntimeError as e: 1264 if len(e.args) > 0 and "recursion" in e.args[0]:

TypeError: cannot pickle 'generator' object

#### Versions:
> cellrank==2.0.4 scanpy==1.10.2 anndata==0.10.8 numpy==1.26.4 numba==0.60.0 scipy==1.11.4 pandas==2.2.2 pygpcca==1.0.4 scikit-learn==1.0.2 statsmodels==0.13.2 scvelo==0.3.2 pygam==0.9.1 matplotlib==3.9.1 seaborn==0.13.2

[eroell]

We see a similar error with the compute_lineage_drivers. Interestingly, if running a juptyer notebook in a few attempts, at one point this error disappears and it works. There is something up with the caching by numba I believe.

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/caching.py:469, in IndexDataCacheFile.save(self, key, data)
    467 try:
    468     # If key already exists, we will overwrite the file
--> 469     data_name = overloads[key]
    470 except KeyError:
    471     # Find an available name for the data file
KeyError: ((Function(<function mean at 0x7f52346bbd80>), int64, Array(float32, 2, 'F', False, aligned=True)), ('x86_64-unknown-linux-gnu', 'znver3', '+64bit,+adx,+aes,-amx-bf16,-amx-int8,-amx-tile,+avx,+avx2,-avx512bf16,-avx512bitalg,-avx512bw,-avx512cd,-avx512dq,-avx512er,-avx512f,-avx512fp16,-avx512ifma,-avx512pf,-avx512vbmi,-avx512vbmi2,-avx512vl,-avx512vnni,-avx512vp2intersect,-avx512vpopcntdq,-avxvnni,+bmi,+bmi2,-cldemote,+clflushopt,+clwb,+clzero,+cmov,+crc32,+cx16,+cx8,-enqcmd,+f16c,+fma,-fma4,+fsgsbase,+fxsr,-gfni,-hreset,+invpcid,-kl,-lwp,+lzcnt,+mmx,+movbe,-movdir64b,-movdiri,-mwaitx,+pclmul,-pconfig,-pku,+popcnt,-prefetchwt1,+prfchw,-ptwrite,+rdpid,+rdrnd,+rdseed,-rtm,+sahf,-serialize,-sgx,+sha,+shstk,+sse,+sse2,+sse3,+sse4.1,+sse4.2,+sse4a,+ssse3,-tbm,-tsxldtrk,-uintr,+vaes,+vpclmulqdq,-waitpkg,-wbnoinvd,-widekl,-xop,+xsave,+xsavec,+xsaveopt,+xsaves'), ('7fab28091f3a48e3ee7271079cbcd8c17f1fe60a518c1d7361068c4d7e3792b0', 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855')
During handling of the above exception, another exception occurred:
TypeError                                 Traceback (most recent call last)
Cell In[22], line 1
----> 1 drivers_1_1 = g.compute_lineage_drivers(lineages="1_1")
      2 adata.obs["fate_probs_1_1"] = g.fate_probabilities["1_1"].X.flatten()
      4 ep.pl.umap(
      5     adata,
      6     color=["fate_probs_1_1"] + list(drivers_1_1.index[:8]),
   (...)
     10     vmax="p96",
     11 )
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/cellrank/estimators/mixins/_lineage_drivers.py:189, in LinDriversMixin.compute_lineage_drivers(self, lineages, method, cluster_key, clusters, layer, use_raw, confidence_level, n_perms, seed, **kwargs)
    183 start = logg.debug(
    184     f"Computing correlations for lineages `{sorted(lineages)}` restricted to clusters `{clusters}` in "
    185     f"layer `{'X' if layer is None else layer}` with `use_raw={use_raw}`"
    186 )
    188 lin_probs = lin_probs[lineages]
--> 189 drivers = _correlation_test(
    190     data,
    191     lin_probs,
    192     gene_names=var_names,
    193     method=method,
    194     n_perms=n_perms,
    195     seed=seed,
    196     confidence_level=confidence_level,
    197     **kwargs,
    198 )
    199 params = self._create_params()
    200 self._write_lineage_drivers(drivers.loc[var_names], use_raw=use_raw, params=params, time=start)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/cellrank/_utils/_utils.py:405, in _correlation_test(X, Y, gene_names, method, confidence_level, n_perms, seed, **kwargs)
    361 @d.get_sections(base="correlation_test", sections=["Returns"])
    362 @d.dedent
    363 def _correlation_test(
   (...)
    371     **kwargs: Any,
    372 ) -> pd.DataFrame:
    373     """Perform a statistical test.
    374 
    375     Return NaN for genes which don't vary across cells.
   (...)
    403     - ``'{lineage}_ci_high'`` - upper bound of the ``confidence_level`` correlation confidence interval.
    404     """
--> 405     corr, pvals, ci_low, ci_high = _correlation_test_helper(
    406         X.T,
    407         Y.X,
    408         method=method,
    409         n_perms=n_perms,
    410         seed=seed,
    411         confidence_level=confidence_level,
    412         **kwargs,
    413     )
    414     invalid = (corr < -1) | (corr > 1)
    415     if np.any(invalid):
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/cellrank/_utils/_utils.py:498, in _correlation_test_helper(X, Y, method, n_perms, seed, confidence_level, **kwargs)
    495 if sp.issparse(X) and not sp.isspmatrix_csr(X):
    496     X = sp.csr_matrix(X)
--> 498 corr = _mat_mat_corr_sparse(X, Y) if sp.issparse(X) else _mat_mat_corr_dense(X, Y)
    500 if method == TestMethod.FISHER:
    501     # see: https://en.wikipedia.org/wiki/Pearson_correlation_coefficient#Using_the_Fisher_transformation
    502     mean, se = np.arctanh(corr), 1.0 / np.sqrt(n - 3)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/cellrank/_utils/_utils.py:318, in _mat_mat_corr_dense(X, Y)
    314 from cellrank.kernels._utils import np_mean, np_std
    316 n = X.shape[1]
--> 318 X_bar = np.reshape(np_mean(X, axis=1), (-1, 1))
    319 X_std = np.reshape(np_std(X, axis=1), (-1, 1))
    321 y_bar = np.reshape(np_mean(Y, axis=0), (1, -1))
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:442, in _DispatcherBase._compile_for_args(self, *args, **kws)
    440             e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
    441     # ignore the FULL_TRACEBACKS config, this needs reporting!
--> 442     raise e
    443 finally:
    444     self._types_active_call = []
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:375, in _DispatcherBase._compile_for_args(self, *args, **kws)
    373 return_val = None
    374 try:
--> 375     return_val = self.compile(tuple(argtypes))
    376 except errors.ForceLiteralArg as e:
    377     # Received request for compiler re-entry with the list of arguments
    378     # indicated by e.requested_args.
    379     # First, check if any of these args are already Literal-ized
    380     already_lit_pos = [i for i in e.requested_args
    381                        if isinstance(args[i], types.Literal)]
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:905, in Dispatcher.compile(self, sig)
    903 with ev.trigger_event("numba:compile", data=ev_details):
    904     try:
--> 905         cres = self._compiler.compile(args, return_type)
    906     except errors.ForceLiteralArg as e:
    907         def folded(args, kws):
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:80, in _FunctionCompiler.compile(self, args, return_type)
     79 def compile(self, args, return_type):
---> 80     status, retval = self._compile_cached(args, return_type)
     81     if status:
     82         return retval
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:94, in _FunctionCompiler._compile_cached(self, args, return_type)
     91     pass
     93 try:
---> 94     retval = self._compile_core(args, return_type)
     95 except errors.TypingError as e:
     96     self._failed_cache[key] = e
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:107, in _FunctionCompiler._compile_core(self, args, return_type)
    104 flags = self._customize_flags(flags)
    106 impl = self._get_implementation(args, {})
--> 107 cres = compiler.compile_extra(self.targetdescr.typing_context,
    108                               self.targetdescr.target_context,
    109                               impl,
    110                               args=args, return_type=return_type,
    111                               flags=flags, locals=self.locals,
    112                               pipeline_class=self.pipeline_class)
    113 # Check typing error if object mode is used
    114 if cres.typing_error is not None and not flags.enable_pyobject:
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler.py:744, in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
    720 """Compiler entry point
    721 
    722 Parameter
   (...)
    740     compiler pipeline
    741 """
    742 pipeline = pipeline_class(typingctx, targetctx, library,
    743                           args, return_type, flags, locals)
--> 744 return pipeline.compile_extra(func)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler.py:438, in CompilerBase.compile_extra(self, func)
    436 self.state.lifted = ()
    437 self.state.lifted_from = None
--> 438 return self._compile_bytecode()
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler.py:506, in CompilerBase._compile_bytecode(self)
    502 """
    503 Populate and run pipeline for bytecode input
    504 """
    505 assert self.state.func_ir is None
--> 506 return self._compile_core()
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler.py:481, in CompilerBase._compile_core(self)
    478 except Exception as e:
    479     if (utils.use_new_style_errors() and not
    480             isinstance(e, errors.NumbaError)):
--> 481         raise e
    483     self.state.status.fail_reason = e
    484     if is_final_pipeline:
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler.py:472, in CompilerBase._compile_core(self)
    470 res = None
    471 try:
--> 472     pm.run(self.state)
    473     if self.state.cr is not None:
    474         break
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler_machinery.py:364, in PassManager.run(self, state)
    361 except Exception as e:
    362     if (utils.use_new_style_errors() and not
    363             isinstance(e, errors.NumbaError)):
--> 364         raise e
    365     msg = "Failed in %s mode pipeline (step: %s)" % \
    366         (self.pipeline_name, pass_desc)
    367     patched_exception = self._patch_error(msg, e)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler_machinery.py:356, in PassManager.run(self, state)
    354 pass_inst = _pass_registry.get(pss).pass_inst
    355 if isinstance(pass_inst, CompilerPass):
--> 356     self._runPass(idx, pass_inst, state)
    357 else:
    358     raise BaseException("Legacy pass in use")
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler_lock.py:35, in _CompilerLock.__call__.<locals>._acquire_compile_lock(*args, **kwargs)
     32 @functools.wraps(func)
     33 def _acquire_compile_lock(*args, **kwargs):
     34     with self:
---> 35         return func(*args, **kwargs)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler_machinery.py:311, in PassManager._runPass(self, index, pss, internal_state)
    309     mutated |= check(pss.run_initialization, internal_state)
    310 with SimpleTimer() as pass_time:
--> 311     mutated |= check(pss.run_pass, internal_state)
    312 with SimpleTimer() as finalize_time:
    313     mutated |= check(pss.run_finalizer, internal_state)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/compiler_machinery.py:273, in PassManager._runPass.<locals>.check(func, compiler_state)
    272 def check(func, compiler_state):
--> 273     mangled = func(compiler_state)
    274     if mangled not in (True, False):
    275         msg = ("CompilerPass implementations should return True/False. "
    276                "CompilerPass with name '%s' did not.")
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typed_passes.py:112, in BaseTypeInference.run_pass(self, state)
    106 """
    107 Type inference and legalization
    108 """
    109 with fallback_context(state, 'Function "%s" failed type inference'
    110                       % (state.func_id.func_name,)):
    111     # Type inference
--> 112     typemap, return_type, calltypes, errs = type_inference_stage(
    113         state.typingctx,
    114         state.targetctx,
    115         state.func_ir,
    116         state.args,
    117         state.return_type,
    118         state.locals,
    119         raise_errors=self._raise_errors)
    120     state.typemap = typemap
    121     # save errors in case of partial typing
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typed_passes.py:93, in type_inference_stage(typingctx, targetctx, interp, args, return_type, locals, raise_errors)
     91     infer.build_constraint()
     92     # return errors in case of partial typing
---> 93     errs = infer.propagate(raise_errors=raise_errors)
     94     typemap, restype, calltypes = infer.unify(raise_errors=raise_errors)
     96 return _TypingResults(typemap, restype, calltypes, errs)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typeinfer.py:1083, in TypeInferer.propagate(self, raise_errors)
   1080 oldtoken = newtoken
   1081 # Errors can appear when the type set is incomplete; only
   1082 # raise them when there is no progress anymore.
-> 1083 errors = self.constraints.propagate(self)
   1084 newtoken = self.get_state_token()
   1085 self.debug.propagate_finished()
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typeinfer.py:182, in ConstraintNetwork.propagate(self, typeinfer)
    180     errors.append(utils.chain_exception(new_exc, e))
    181 elif utils.use_new_style_errors():
--> 182     raise e
    183 else:
    184     msg = ("Unknown CAPTURED_ERRORS style: "
    185            f"'{config.CAPTURED_ERRORS}'.")
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typeinfer.py:160, in ConstraintNetwork.propagate(self, typeinfer)
    157 with typeinfer.warnings.catch_warnings(filename=loc.filename,
    158                                        lineno=loc.line):
    159     try:
--> 160         constraint(typeinfer)
    161     except ForceLiteralArg as e:
    162         errors.append(e)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typeinfer.py:583, in CallConstraint.__call__(self, typeinfer)
    581     fnty = typevars[self.func].getone()
    582 with new_error_context("resolving callee type: {0}", fnty):
--> 583     self.resolve(typeinfer, typevars, fnty)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typeinfer.py:606, in CallConstraint.resolve(self, typeinfer, typevars, fnty)
    604     fnty = fnty.instance_type
    605 try:
--> 606     sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
    607 except ForceLiteralArg as e:
    608     # Adjust for bound methods
    609     folding_args = ((fnty.this,) + tuple(self.args)
    610                     if isinstance(fnty, types.BoundFunction)
    611                     else self.args)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typeinfer.py:1577, in TypeInferer.resolve_call(self, fnty, pos_args, kw_args)
   1574     return sig
   1575 else:
   1576     # Normal non-recursive call
-> 1577     return self.context.resolve_function_type(fnty, pos_args, kw_args)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typing/context.py:196, in BaseContext.resolve_function_type(self, func, args, kws)
    194 # Prefer user definition first
    195 try:
--> 196     res = self._resolve_user_function_type(func, args, kws)
    197 except errors.TypingError as e:
    198     # Capture any typing error
    199     last_exception = e
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/typing/context.py:248, in BaseContext._resolve_user_function_type(self, func, args, kws, literals)
    244         return self.resolve_function_type(func_type, args, kws)
    246 if isinstance(func, types.Callable):
    247     # XXX fold this into the __call__ attribute logic?
--> 248     return func.get_call_type(self, args, kws)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/types/functions.py:541, in Dispatcher.get_call_type(self, context, args, kws)
    534 def get_call_type(self, context, args, kws):
    535     """
    536     Resolve a call to this dispatcher using the given argument types.
    537     A signature returned and it is ensured that a compiled specialization
    538     is available for it.
    539     """
    540     template, pysig, args, kws = \
--> 541         self.dispatcher.get_call_template(args, kws)
    542     sig = template(context).apply(args, kws)
    543     if sig:
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:318, in _DispatcherBase.get_call_template(self, args, kws)
    316 # Ensure an overload is available
    317 if self._can_compile:
--> 318     self.compile(tuple(args))
    320 # Create function type for typing
    321 func_name = self.py_func.__name__
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/dispatcher.py:912, in Dispatcher.compile(self, sig)
    910         raise e.bind_fold_arguments(folded)
    911     self.add_overload(cres)
--> 912 self._cache.save_overload(sig, cres)
    913 return cres.entry_point
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/caching.py:652, in Cache.save_overload(self, sig, data)
    648 """
    649 Save the data for the given signature in the cache.
    650 """
    651 with self._guard_against_spurious_io_errors():
--> 652     self._save_overload(sig, data)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/caching.py:662, in Cache._save_overload(self, sig, data)
    660 key = self._index_key(sig, data.codegen)
    661 data = self._impl.reduce(data)
--> 662 self._cache_file.save(key, data)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/caching.py:478, in IndexDataCacheFile.save(self, key, data)
    476             break
    477     overloads[key] = data_name
--> 478     self._save_index(overloads)
    479 self._save_data(data_name, data)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/caching.py:522, in IndexDataCacheFile._save_index(self, overloads)
    520 def _save_index(self, overloads):
    521     data = self._source_stamp, overloads
--> 522     data = self._dump(data)
    523     with self._open_for_write(self._index_path) as f:
    524         pickle.dump(self._version, f, protocol=-1)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/caching.py:550, in IndexDataCacheFile._dump(self, obj)
    549 def _dump(self, obj):
--> 550     return dumps(obj)
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/core/serialize.py:58, in dumps(obj)
     56 with io.BytesIO() as buf:
     57     p = pickler(buf, protocol=4)
---> 58     p.dump(obj)
     59     pickled = buf.getvalue()
     61 return pickled
File /opt/hostedtoolcache/Python/3.11.9/x64/lib/python3.11/site-packages/numba/cloudpickle/cloudpickle.py:1262, in Pickler.dump(self, obj)
   1260 def dump(self, obj):
   1261     try:
-> 1262         return super().dump(obj)
   1263     except RuntimeError as e:
   1264         if len(e.args) > 0 and "recursion" in e.args[0]:
TypeError: cannot pickle 'generator' object
Error: Process completed with exit code 1.

[Marius1311]

@michalk8, any idea of what might be going on there?

[TTTPOB]

me too, after hit enter for 5 times, it runs.

[dsb66]

I have the same problem. Using the trick of hitting run 5 times, I got the computation to start and complete the first progress bar, but then execution fails at the second progress bar with this error message:

00%
 30763/30763 [00:14<00:00, 2111.47cell/s]
/opt/bin/miniforge3/envs/scanpy/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=179552) is multi-threaded, use of fork() may lead to deadlocks in the child.
  self.pid = os.fork()
  0%
 0/30763 [00:00<?, ?cell/s]
/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/externals/loky/backend/fork_exec.py:38: DeprecationWarning: This process (pid=179552) is multi-threaded, use of fork() may lead to deadlocks in the child.
  pid = os.fork()
---------------------------------------------------------------------------
_RemoteTraceback                          Traceback (most recent call last)
_RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/caching.py", line 469, in save
    data_name = overloads[key]
                ~~~~~~~~~^^^^^
KeyError: ((Function(<function mean at 0x7ffff34ac2c0>), Literal[int](1), Array(float64, 2, 'A', True, aligned=True)), ('x86_64-unknown-linux-gnu', 'icelake-server', '+64bit,+adx,+aes,-amx-bf16,-amx-int8,-amx-tile,+avx,+avx2,-avx512bf16,+avx512bitalg,+avx512bw,+avx512cd,+avx512dq,-avx512er,+avx512f,-avx512fp16,+avx512ifma,-avx512pf,+avx512vbmi,+avx512vbmi2,+avx512vl,+avx512vnni,-avx512vp2intersect,+avx512vpopcntdq,-avxvnni,+bmi,+bmi2,-cldemote,+clflushopt,+clwb,-clzero,+cmov,+crc32,+cx16,+cx8,-enqcmd,+f16c,+fma,-fma4,+fsgsbase,+fxsr,+gfni,-hreset,+invpcid,-kl,-lwp,+lzcnt,+mmx,+movbe,-movdir64b,-movdiri,-mwaitx,+pclmul,+pconfig,+pku,+popcnt,-prefetchwt1,+prfchw,-ptwrite,+rdpid,+rdrnd,+rdseed,-rtm,+sahf,-serialize,+sgx,+sha,-shstk,+sse,+sse2,+sse3,+sse4.1,+sse4.2,-sse4a,+ssse3,-tbm,-tsxldtrk,-uintr,+vaes,+vpclmulqdq,-waitpkg,+wbnoinvd,-widekl,-xop,+xsave,+xsavec,+xsaveopt,+xsaves'), ('d2da2eff5021c157013871a5792da29a8ae93428b6fe2f2daf7c8c3f3c1c25ba', 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'))

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/externals/loky/process_executor.py", line 463, in _process_worker
    r = call_item()
        ^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/externals/loky/process_executor.py", line 291, in __call__
    return self.fn(*self.args, **self.kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py", line 598, in __call__
    return [func(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/cellrank/kernels/utils/_velocity_model.py", line 85, in _compute_helper
    ps, ls = self._compute(ix, neigh_ixs, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/cellrank/kernels/utils/_velocity_model.py", line 181, in _compute
    return self._similarity(v[None, :], W, self._softmax_scale)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/cellrank/kernels/utils/_similarity.py", line 238, in __call__
    return _predict_transition_probabilities_numpy(v, D, softmax_scale, self._center_mean, self._scale_by_norm)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 442, in _compile_for_args
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 375, in _compile_for_args
    return_val = self.compile(tuple(argtypes))
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 905, in compile
    cres = self._compiler.compile(args, return_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 80, in compile
    status, retval = self._compile_cached(args, return_type)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 94, in _compile_cached
    retval = self._compile_core(args, return_type)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 107, in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 744, in compile_extra
    return pipeline.compile_extra(func)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 438, in compile_extra
    return self._compile_bytecode()
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 506, in _compile_bytecode
    return self._compile_core()
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 481, in _compile_core
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 472, in _compile_core
    pm.run(self.state)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 364, in run
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 356, in run
    self._runPass(idx, pass_inst, state)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 311, in _runPass
    mutated |= check(pss.run_pass, internal_state)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 273, in check
    mangled = func(compiler_state)
              ^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typed_passes.py", line 112, in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(
                                            ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typed_passes.py", line 93, in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 1083, in propagate
    errors = self.constraints.propagate(self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 182, in propagate
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 160, in propagate
    constraint(typeinfer)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 583, in __call__
    self.resolve(typeinfer, typevars, fnty)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 606, in resolve
    sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 1577, in resolve_call
    return self.context.resolve_function_type(fnty, pos_args, kw_args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typing/context.py", line 196, in resolve_function_type
    res = self._resolve_user_function_type(func, args, kws)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typing/context.py", line 248, in _resolve_user_function_type
    return func.get_call_type(self, args, kws)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/types/functions.py", line 541, in get_call_type
    self.dispatcher.get_call_template(args, kws)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 318, in get_call_template
    self.compile(tuple(args))
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 905, in compile
    cres = self._compiler.compile(args, return_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 80, in compile
    status, retval = self._compile_cached(args, return_type)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 94, in _compile_cached
    retval = self._compile_core(args, return_type)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 107, in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 744, in compile_extra
    return pipeline.compile_extra(func)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 438, in compile_extra
    return self._compile_bytecode()
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 506, in _compile_bytecode
    return self._compile_core()
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 481, in _compile_core
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler.py", line 472, in _compile_core
    pm.run(self.state)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 364, in run
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 356, in run
    self._runPass(idx, pass_inst, state)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 311, in _runPass
    mutated |= check(pss.run_pass, internal_state)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/compiler_machinery.py", line 273, in check
    mangled = func(compiler_state)
              ^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typed_passes.py", line 112, in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(
                                            ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typed_passes.py", line 93, in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 1083, in propagate
    errors = self.constraints.propagate(self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 182, in propagate
    raise e
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 160, in propagate
    constraint(typeinfer)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 583, in __call__
    self.resolve(typeinfer, typevars, fnty)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 606, in resolve
    sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typeinfer.py", line 1577, in resolve_call
    return self.context.resolve_function_type(fnty, pos_args, kw_args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typing/context.py", line 196, in resolve_function_type
    res = self._resolve_user_function_type(func, args, kws)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/typing/context.py", line 248, in _resolve_user_function_type
    return func.get_call_type(self, args, kws)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/types/functions.py", line 541, in get_call_type
    self.dispatcher.get_call_template(args, kws)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 318, in get_call_template
    self.compile(tuple(args))
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/dispatcher.py", line 912, in compile
    self._cache.save_overload(sig, cres)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/caching.py", line 652, in save_overload
    self._save_overload(sig, data)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/caching.py", line 662, in _save_overload
    self._cache_file.save(key, data)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/caching.py", line 478, in save
    self._save_index(overloads)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/caching.py", line 522, in _save_index
    data = self._dump(data)
           ^^^^^^^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/caching.py", line 550, in _dump
    return dumps(obj)
           ^^^^^^^^^^
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/core/serialize.py", line 58, in dumps
    p.dump(obj)
  File "/opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/numba/cloudpickle/cloudpickle.py", line 1262, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'generator' object
"""

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

TypeError                                 Traceback (most recent call last)
Cell In[20], line 3
      1 # velocity kernel, transition probs
      2 vk = cr.kernels.VelocityKernel(adata)
----> 3 vk.compute_transition_matrix(n_jobs=ncpu)

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/cellrank/kernels/_velocity_kernel.py:175, in VelocityKernel.compute_transition_matrix(self, model, backward_mode, similarity, softmax_scale, n_samples, seed, **kwargs)
    173 if isinstance(model, Stochastic):
    174     kwargs["backend"] = DEFAULT_BACKEND
--> 175 self.transition_matrix, self._logits = model(**kwargs)
    177 logg.info("    Finish", time=start)
    179 return self

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/cellrank/kernels/utils/_velocity_model.py:64, in ModelABC.__call__(self, n_jobs, backend, show_progress_bar, **kwargs)
     56 def __call__(
     57     self,
     58     n_jobs: Optional[int] = None,
   (...)
     61     **kwargs: Any,
     62 ) -> Tuple[np.ndarray, np.ndarray]:
     63     ixs = self._ixs
---> 64     return parallelize(
     65         self._compute_helper,
     66         ixs,
     67         n_jobs=n_jobs,
     68         backend=backend,
     69         show_progress_bar=show_progress_bar,
     70         as_array=False,
     71         extractor=lambda data: self._reconstruct_output(np.concatenate(data, axis=-1), ixs),
     72         unit=self._unit,
     73     )(**kwargs)

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/cellrank/_utils/_parallelize.py:96, in parallelize.<locals>.wrapper(*args, **kwargs)
     93 else:
     94     pbar, queue, thread = None, None, None
---> 96 res = jl.Parallel(n_jobs=n_jobs, backend=backend)(
     97     jl.delayed(callback)(
     98         *((i, cs) if use_ixs else (cs,)),
     99         *args,
    100         **kwargs,
    101         queue=queue,
    102     )
    103     for i, cs in enumerate(collections)
    104 )
    106 res = np.array(res) if as_array else res
    107 if thread is not None:

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py:2007, in Parallel.__call__(self, iterable)
   2001 # The first item from the output is blank, but it makes the interpreter
   2002 # progress until it enters the Try/Except block of the generator and
   2003 # reaches the first `yield` statement. This starts the asynchronous
   2004 # dispatch of the tasks to the workers.
   2005 next(output)
-> 2007 return output if self.return_generator else list(output)

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py:1650, in Parallel._get_outputs(self, iterator, pre_dispatch)
   1647     yield
   1649     with self._backend.retrieval_context():
-> 1650         yield from self._retrieve()
   1652 except GeneratorExit:
   1653     # The generator has been garbage collected before being fully
   1654     # consumed. This aborts the remaining tasks if possible and warn
   1655     # the user if necessary.
   1656     self._exception = True

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py:1754, in Parallel._retrieve(self)
   1747 while self._wait_retrieval():
   1748 
   1749     # If the callback thread of a worker has signaled that its task
   1750     # triggered an exception, or if the retrieval loop has raised an
   1751     # exception (e.g. `GeneratorExit`), exit the loop and surface the
   1752     # worker traceback.
   1753     if self._aborting:
-> 1754         self._raise_error_fast()
   1755         break
   1757     # If the next job is not ready for retrieval yet, we just wait for
   1758     # async callbacks to progress.

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py:1789, in Parallel._raise_error_fast(self)
   1785 # If this error job exists, immediately raise the error by
   1786 # calling get_result. This job might not exists if abort has been
   1787 # called directly or if the generator is gc'ed.
   1788 if error_job is not None:
-> 1789     error_job.get_result(self.timeout)

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py:745, in BatchCompletionCallBack.get_result(self, timeout)
    739 backend = self.parallel._backend
    741 if backend.supports_retrieve_callback:
    742     # We assume that the result has already been retrieved by the
    743     # callback thread, and is stored internally. It's just waiting to
    744     # be returned.
--> 745     return self._return_or_raise()
    747 # For other backends, the main thread needs to run the retrieval step.
    748 try:

File /opt/bin/miniforge3/envs/scanpy/lib/python3.12/site-packages/joblib/parallel.py:763, in BatchCompletionCallBack._return_or_raise(self)
    761 try:
    762     if self.status == TASK_ERROR:
--> 763         raise self._result
    764     return self._result
    765 finally:

TypeError: cannot pickle 'generator' object

[dsb66]

The issue appears to be with parallelization. If I run the command as vk.compute_transition_matrix(n_jobs=1, show_progress_bar=False) it completes successfully. I am running cellrank in jupyter on an HPC running slurm. When I run the command with n_jobs equal to the number of CPUs assigned to my session, the command fails with the error message above.

[michalk8]

Fixed in #1212, will create a new release shortly.

[Marius1311]

Amazing, thanks a lot @michalk8!