Multiprocess crash

Hi, my simulation program frequently crashed during the multiprocess while it works during single thread process. I used multiprocessing package and tested python v 3.11, v 3.12 along with pymc 5.10.2 and 5.9.2 but didn’t work out (using conda-forge installation and tested 2x2=4 situations). Does anyone has experienced same thing?

main part of my code

def single_round_simulation(true_beta, RANDOM_SEED, samples_size, chain_len):
    rng = np.random.default_rng(RANDOM_SEED)
    # skip some code because of it is too long
    basic_model = pm.Model()
    with basic_model:
        # similar to the user guide code
    with basic_model:
        idata = pm.sample(draws=chain_len, chains=1)
    return res
      
ncores = 24 
nrepeat= 1000
true_beta_l = [0.2, 0.2, 0.3, 0.3]
# true_beta_l = [0., 0., 0., 0.]
# RANDOM_SEED = 24
samples_size = 1000
chain_len = 10000
mcmc_betas_mean_l, mcmc_betas_std_l, ols_betas_mean_l, ols_betas_std_l = [], [], [], []

start_time = time.time()
with mp.Pool(ncores) as pool:
    temp_reses = [ # return 
        pool.apply_async(
            single_round_simulation, # function
            (true_beta_l, RANDOM_SEED, samples_size, chain_len),
        )
        for RANDOM_SEED in range(nrepeat) # 随机种子
    ]
    for temp_resi in tqdm(temp_reses, desc="Pipeline: "):
        ( # 函数的return
            mcmc_betas_mean, 
            mcmc_betas_std, 
            ols_betas_mean, 
            ols_betas_std
        ) = temp_resi.get()
        mcmc_betas_mean_l.append(mcmc_betas_mean)
        mcmc_betas_std_l.append(mcmc_betas_std)
        ols_betas_mean_l.append(ols_betas_mean)
        ols_betas_std_l.append(ols_betas_std)

mcmc_betas_mean_l = np.array(mcmc_betas_mean_l)
mcmc_betas_std_l = np.array(mcmc_betas_std_l)
ols_betas_mean_l = np.array(ols_betas_mean_l)
ols_betas_std_l = np.array(ols_betas_std_l)
end_time = time.time()

the error log

Pipeline:   0%|          | 2/1000 [04:34<38:04:46, 137.36s/it]
multiprocessing.pool.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/vm.py", line 1235, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 119, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 84, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1209, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1129, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1653, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1197, in module_from_key
    module = self._get_from_hash(module_hash, key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1111, in _get_from_hash
    self.check_key(key, key_data.key_pkl)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1306, in check_key
    raise AssertionError(
AssertionError: Key not found in unpickled KeyData file. Verify the __eq__ and __hash__ functions of your Ops. The file is: /home/user/.pytensor/compiledir_Linux-3.10-el7.x86_64-x86_64-with-glibc2.17-x86_64-3.11.7-64/tmphigv8k2c/key.pkl. The key is: (((15, (4, (4,), (4,), (4,), (4,), (4,), (4,), (4,), (4,), (4,)), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), ('openmp', False), ('openmp_elemwise_minsize', 200000)), ('scalar_op', 'inplace_pattern'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2')), ('CLinker.cmodule_key', ('--param', '--param', '--param', '-DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION', '-O3', '-Wno-c++11-narrowing', '-Wno-unused-label', '-Wno-unused-variable', '-Wno-write-strings', '-fPIC', '-fno-asynchronous-unwind-tables', '-fno-exceptions', '-fno-math-errno', '-fno-unwind-tables', '-m64', '-mabm', '-madx', '-maes', '-march=cascadelake', '-mavx', '-mavx2', '-mavx512bw', '-mavx512cd', '-mavx512dq', '-mavx512f', '-mavx512vl', '-mavx512vnni', '-mbmi', '-mbmi2', '-mclflushopt', '-mclwb', '-mcx16', '-mf16c', '-mfma', '-mfsgsbase', '-mfxsr', '-mhle', '-mlzcnt', '-mmmx', '-mmovbe', '-mno-3dnow', '-mno-amx-bf16', '-mno-amx-int8', '-mno-amx-tile', '-mno-avx5124fmaps', '-mno-avx5124vnniw', '-mno-avx512bf16', '-mno-avx512bitalg', '-mno-avx512er', '-mno-avx512fp16', '-mno-avx512ifma', '-mno-avx512pf', '-mno-avx512vbmi', '-mno-avx512vbmi2', '-mno-avx512vp2intersect', '-mno-avx512vpopcntdq', '-mno-avxvnni', '-mno-cldemote', '-mno-clzero', '-mno-enqcmd', '-mno-fma4', '-mno-gfni', '-mno-hreset', '-mno-kl', '-mno-lwp', '-mno-movdir64b', '-mno-movdiri', '-mno-mwaitx', '-mno-pconfig', '-mno-prefetchwt1', '-mno-ptwrite', '-mno-rdpid', '-mno-serialize', '-mno-sgx', '-mno-sha', '-mno-shstk', '-mno-sse4a', '-mno-tbm', '-mno-tsxldtrk', '-mno-uintr', '-mno-vaes', '-mno-vpclmulqdq', '-mno-waitpkg', '-mno-wbnoinvd', '-mno-widekl', '-mno-xop', '-mpclmul', '-mpku', '-mpopcnt', '-mprfchw', '-mrdrnd', '-mrdseed', '-mrtm', '-msahf', '-msse', '-msse2', '-msse3', '-msse4.1', '-msse4.2', '-mssse3', '-mtune=cascadelake', '-mxsave', '-mxsavec', '-mxsaveopt', '-mxsaves', 'l1-cache-line-size=64', 'l1-cache-size=32', 'l2-cache-size=36608'), (), (), 'NPY_ABI_VERSION=0x1000009', 'c_compiler_str=/home/user/miniforge3/envs/pymc_env0/bin/g++ 12.3.0', 'md5:m624562459f98c9fea8456c836a7c615e5d233a9e9905eebffeb353cd3b5f9076', (Elemwise(scalar_op=Composite{...},inplace_pattern=<frozendict {}>), ((TensorType(float64, shape=(1000,)), ((-1, 0), False)), (TensorType(float64, shape=(1,)), ((-1, 1), False)), (TensorType(float64, shape=(1000,)), (('mc6bff54bc00ec2145cb60946612bafac978a4c3452f0851b4829d349d7c7fb98', 0, 2), False)), (TensorType(float64, shape=(1000,)), (('m3f433ed8e956d44b2e773ea64de398180ca995df3d9df3db4fb4c722d39d8e80', 0, 3), False)), (TensorType(float64, shape=(1,)), ((-1, 4), False)), (TensorType(bool, shape=(1,)), ((-1, 5), False))), (1, (False, False, False, False)))))

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
                    ^^^^^^^^^^^^^^^^^^^
  File "/home/user/rwd_study/mp.py", line 74, in single_round_simulation
    idata = pm.sample(draws=chain_len, chains=1)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 689, in sample
    step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 239, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 140, in instantiate_steppers
    step = step_class(vars=vars, model=model, **args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/hmc/nuts.py", line 180, in __init__
    super().__init__(vars, **kwargs)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/hmc/base_hmc.py", line 109, in __init__
    super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **pytensor_kwargs)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/arraystep.py", line 164, in __init__
    func = model.logp_dlogp_function(vars, dtype=dtype, **pytensor_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/model/core.py", line 618, in logp_dlogp_function
    return ValueGradFunction(costs, grad_vars, extra_vars_and_values, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/model/core.py", line 350, in __init__
    self._pytensor_function = compile_pymc(inputs, outputs, givens=givens, **kwargs)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/pytensorf.py", line 991, in compile_pymc
    pytensor_function = pytensor.function(
                        ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/__init__.py", line 315, in function
    fn = pfunc(
         ^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py", line 469, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1762, in orig_function
    fn = m.create(defaults)
         ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1654, in create
    _fn, _i, _o = self.linker.make_thunk(
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/basic.py", line 245, in make_thunk
    return self.make_all(
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/vm.py", line 1244, in make_all
    raise_with_op(fgraph, node)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/utils.py", line 531, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/vm.py", line 1235, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 119, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 84, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1209, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1129, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1653, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1197, in module_from_key
    module = self._get_from_hash(module_hash, key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1111, in _get_from_hash
    self.check_key(key, key_data.key_pkl)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1306, in check_key
    raise AssertionError(
AssertionError: Key not found in unpickled KeyData file. Verify the __eq__ and __hash__ functions of your Ops. The file is: /home/user/.pytensor/compiledir_Linux-3.10-el7.x86_64-x86_64-with-glibc2.17-x86_64-3.11.7-64/tmphigv8k2c/key.pkl. The key is: (((15, (4, (4,), (4,), (4,), (4,), (4,), (4,), (4,), (4,), (4,)), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), ('openmp', False), ('openmp_elemwise_minsize', 200000)), ('scalar_op', 'inplace_pattern'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2')), ('CLinker.cmodule_key', ('--param', '--param', '--param', '-DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION', '-O3', '-Wno-c++11-narrowing', '-Wno-unused-label', '-Wno-unused-variable', '-Wno-write-strings', '-fPIC', '-fno-asynchronous-unwind-tables', '-fno-exceptions', '-fno-math-errno', '-fno-unwind-tables', '-m64', '-mabm', '-madx', '-maes', '-march=cascadelake', '-mavx', '-mavx2', '-mavx512bw', '-mavx512cd', '-mavx512dq', '-mavx512f', '-mavx512vl', '-mavx512vnni', '-mbmi', '-mbmi2', '-mclflushopt', '-mclwb', '-mcx16', '-mf16c', '-mfma', '-mfsgsbase', '-mfxsr', '-mhle', '-mlzcnt', '-mmmx', '-mmovbe', '-mno-3dnow', '-mno-amx-bf16', '-mno-amx-int8', '-mno-amx-tile', '-mno-avx5124fmaps', '-mno-avx5124vnniw', '-mno-avx512bf16', '-mno-avx512bitalg', '-mno-avx512er', '-mno-avx512fp16', '-mno-avx512ifma', '-mno-avx512pf', '-mno-avx512vbmi', '-mno-avx512vbmi2', '-mno-avx512vp2intersect', '-mno-avx512vpopcntdq', '-mno-avxvnni', '-mno-cldemote', '-mno-clzero', '-mno-enqcmd', '-mno-fma4', '-mno-gfni', '-mno-hreset', '-mno-kl', '-mno-lwp', '-mno-movdir64b', '-mno-movdiri', '-mno-mwaitx', '-mno-pconfig', '-mno-prefetchwt1', '-mno-ptwrite', '-mno-rdpid', '-mno-serialize', '-mno-sgx', '-mno-sha', '-mno-shstk', '-mno-sse4a', '-mno-tbm', '-mno-tsxldtrk', '-mno-uintr', '-mno-vaes', '-mno-vpclmulqdq', '-mno-waitpkg', '-mno-wbnoinvd', '-mno-widekl', '-mno-xop', '-mpclmul', '-mpku', '-mpopcnt', '-mprfchw', '-mrdrnd', '-mrdseed', '-mrtm', '-msahf', '-msse', '-msse2', '-msse3', '-msse4.1', '-msse4.2', '-mssse3', '-mtune=cascadelake', '-mxsave', '-mxsavec', '-mxsaveopt', '-mxsaves', 'l1-cache-line-size=64', 'l1-cache-size=32', 'l2-cache-size=36608'), (), (), 'NPY_ABI_VERSION=0x1000009', 'c_compiler_str=/home/user/miniforge3/envs/pymc_env0/bin/g++ 12.3.0', 'md5:m624562459f98c9fea8456c836a7c615e5d233a9e9905eebffeb353cd3b5f9076', (Elemwise(scalar_op=Composite{...},inplace_pattern=<frozendict {}>), ((TensorType(float64, shape=(1000,)), ((-1, 0), False)), (TensorType(float64, shape=(1,)), ((-1, 1), False)), (TensorType(float64, shape=(1000,)), (('mc6bff54bc00ec2145cb60946612bafac978a4c3452f0851b4829d349d7c7fb98', 0, 2), False)), (TensorType(float64, shape=(1000,)), (('m3f433ed8e956d44b2e773ea64de398180ca995df3d9df3db4fb4c722d39d8e80', 0, 3), False)), (TensorType(float64, shape=(1,)), ((-1, 4), False)), (TensorType(bool, shape=(1,)), ((-1, 5), False))), (1, (False, False, False, False)))))
Apply node that caused the error: Composite{...}(Composite{(i5 - (i2 + (i3 * i4) + (i0 * i1)))}.0, ExpandDims{axis=0}.0, [1. 1. 0. ... 1. 0. 1.], [ 1.111988 ... 68547e+01], Log.0, Gt.0)
Toposort index: 24
Inputs types: [TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1,)), TensorType(bool, shape=(1,))]

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
"""

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

Traceback (most recent call last):
  File "/home/user/rwd_study/mp.py", line 125, in <module>
    ) = temp_resi.get()
        ^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/multiprocessing/pool.py", line 774, in get
    raise self._value
AssertionError: Key not found in unpickled KeyData file. Verify the __eq__ and __hash__ functions of your Ops. The file is: /home/user/.pytensor/compiledir_Linux-3.10-el7.x86_64-x86_64-with-glibc2.17-x86_64-3.11.7-64/tmphigv8k2c/key.pkl. The key is: (((15, (4, (4,), (4,), (4,), (4,), (4,), (4,), (4,), (4,), (4,)), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), (13, '1.26.2'), ('openmp', False), ('openmp_elemwise_minsize', 200000)), ('scalar_op', 'inplace_pattern'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2'), (11, 13, '1.26.2')), ('CLinker.cmodule_key', ('--param', '--param', '--param', '-DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION', '-O3', '-Wno-c++11-narrowing', '-Wno-unused-label', '-Wno-unused-variable', '-Wno-write-strings', '-fPIC', '-fno-asynchronous-unwind-tables', '-fno-exceptions', '-fno-math-errno', '-fno-unwind-tables', '-m64', '-mabm', '-madx', '-maes', '-march=cascadelake', '-mavx', '-mavx2', '-mavx512bw', '-mavx512cd', '-mavx512dq', '-mavx512f', '-mavx512vl', '-mavx512vnni', '-mbmi', '-mbmi2', '-mclflushopt', '-mclwb', '-mcx16', '-mf16c', '-mfma', '-mfsgsbase', '-mfxsr', '-mhle', '-mlzcnt', '-mmmx', '-mmovbe', '-mno-3dnow', '-mno-amx-bf16', '-mno-amx-int8', '-mno-amx-tile', '-mno-avx5124fmaps', '-mno-avx5124vnniw', '-mno-avx512bf16', '-mno-avx512bitalg', '-mno-avx512er', '-mno-avx512fp16', '-mno-avx512ifma', '-mno-avx512pf', '-mno-avx512vbmi', '-mno-avx512vbmi2', '-mno-avx512vp2intersect', '-mno-avx512vpopcntdq', '-mno-avxvnni', '-mno-cldemote', '-mno-clzero', '-mno-enqcmd', '-mno-fma4', '-mno-gfni', '-mno-hreset', '-mno-kl', '-mno-lwp', '-mno-movdir64b', '-mno-movdiri', '-mno-mwaitx', '-mno-pconfig', '-mno-prefetchwt1', '-mno-ptwrite', '-mno-rdpid', '-mno-serialize', '-mno-sgx', '-mno-sha', '-mno-shstk', '-mno-sse4a', '-mno-tbm', '-mno-tsxldtrk', '-mno-uintr', '-mno-vaes', '-mno-vpclmulqdq', '-mno-waitpkg', '-mno-wbnoinvd', '-mno-widekl', '-mno-xop', '-mpclmul', '-mpku', '-mpopcnt', '-mprfchw', '-mrdrnd', '-mrdseed', '-mrtm', '-msahf', '-msse', '-msse2', '-msse3', '-msse4.1', '-msse4.2', '-mssse3', '-mtune=cascadelake', '-mxsave', '-mxsavec', '-mxsaveopt', '-mxsaves', 'l1-cache-line-size=64', 'l1-cache-size=32', 'l2-cache-size=36608'), (), (), 'NPY_ABI_VERSION=0x1000009', 'c_compiler_str=/home/user/miniforge3/envs/pymc_env0/bin/g++ 12.3.0', 'md5:m624562459f98c9fea8456c836a7c615e5d233a9e9905eebffeb353cd3b5f9076', (Elemwise(scalar_op=Composite{...},inplace_pattern=<frozendict {}>), ((TensorType(float64, shape=(1000,)), ((-1, 0), False)), (TensorType(float64, shape=(1,)), ((-1, 1), False)), (TensorType(float64, shape=(1000,)), (('mc6bff54bc00ec2145cb60946612bafac978a4c3452f0851b4829d349d7c7fb98', 0, 2), False)), (TensorType(float64, shape=(1000,)), (('m3f433ed8e956d44b2e773ea64de398180ca995df3d9df3db4fb4c722d39d8e80', 0, 3), False)), (TensorType(float64, shape=(1,)), ((-1, 4), False)), (TensorType(bool, shape=(1,)), ((-1, 5), False))), (1, (False, False, False, False)))))
Apply node that caused the error: Composite{...}(Composite{(i5 - (i2 + (i3 * i4) + (i0 * i1)))}.0, ExpandDims{axis=0}.0, [1. 1. 0. ... 1. 0. 1.], [ 1.111988 ... 68547e+01], Log.0, Gt.0)
Toposort index: 24
Inputs types: [TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1,)), TensorType(bool, shape=(1,))]

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

PyTensor which PyMC uses is limited in what regards multithreading/multiprocessing compilation

Thanks for quick response. Does this mean that currently I can not use multiprocessing in pymc?

Usually users make use of multiprocessing to sample multiple chains in parallel. You can try using the numba backend instead of the C one by setting pytensor.config.change_flags(mode="NUMBA")

Many thanks for the suggestion. I tried the setting and got the folowing error log. Do you have any further suggestion?

---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/vm.py", line 1235, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 119, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 84, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1209, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1129, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1653, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1197, in module_from_key
    module = self._get_from_hash(module_hash, key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1100, in _get_from_hash
    with lock_ctx():
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/contextlib.py", line 137, in __enter__
    return next(self.gen)
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/compilelock.py", line 74, in lock_ctx
    fl.acquire(timeout=timeout)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/filelock/_api.py", line 264, in acquire
    raise Timeout(lock_filename)  # noqa: TRY301
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
filelock._error.Timeout: The file lock '/home/user/.pytensor/compiledir_Linux-3.10-el7.x86_64-x86_64-with-glibc2.17-x86_64-3.11.7-64/.lock' could not be acquired.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
                    ^^^^^^^^^^^^^^^^^^^
  File "/tmp/ipykernel_259166/1255918468.py", line 59, in single_round_simulation
    idata = pm.sample(draws=chain_len, chains=1)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 689, in sample
    step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 239, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 140, in instantiate_steppers
    step = step_class(vars=vars, model=model, **args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/hmc/nuts.py", line 180, in __init__
    super().__init__(vars, **kwargs)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/hmc/base_hmc.py", line 109, in __init__
    super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **pytensor_kwargs)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/arraystep.py", line 164, in __init__
    func = model.logp_dlogp_function(vars, dtype=dtype, **pytensor_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/model/core.py", line 618, in logp_dlogp_function
    return ValueGradFunction(costs, grad_vars, extra_vars_and_values, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/model/core.py", line 350, in __init__
    self._pytensor_function = compile_pymc(inputs, outputs, givens=givens, **kwargs)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/pytensorf.py", line 991, in compile_pymc
    pytensor_function = pytensor.function(
                        ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/__init__.py", line 315, in function
    fn = pfunc(
         ^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py", line 469, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1762, in orig_function
    fn = m.create(defaults)
         ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1654, in create
    _fn, _i, _o = self.linker.make_thunk(
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/basic.py", line 245, in make_thunk
    return self.make_all(
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/vm.py", line 1244, in make_all
    raise_with_op(fgraph, node)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/utils.py", line 531, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/vm.py", line 1235, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 119, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/op.py", line 84, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1209, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1129, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1653, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1197, in module_from_key
    module = self._get_from_hash(module_hash, key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1100, in _get_from_hash
    with lock_ctx():
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/contextlib.py", line 137, in __enter__
    return next(self.gen)
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/compilelock.py", line 74, in lock_ctx
    fl.acquire(timeout=timeout)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/filelock/_api.py", line 264, in acquire
    raise Timeout(lock_filename)  # noqa: TRY301
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
filelock._error.Timeout: The file lock 'The file lock '/home/user/.pytensor/compiledir_Linux-3.10-el7.x86_64-x86_64-with-glibc2.17-x86_64-3.11.7-64/.lock' could not be acquired.
Apply node that caused the error: Composite{(i5 - (i2 + (i3 * i4) + (i0 * i1)))}(ExpandDims{axis=0}.0, [-4.045692 ... 74900e+00], ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, [0. 1. 1. ... 1. 0. 0.], Y_obs1{[-2.025834 ... 00053e+00]})
Toposort index: 6
Inputs types: [TensorType(float64, shape=(1,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1000,))]

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.' could not be acquired.
"""

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

Timeout                                   Traceback (most recent call last)
Cell In[8], line 25
     12 temp_reses = [ # return 
     13     pool.apply_async(
     14         single_round_simulation, # function
   (...)
     17     for RANDOM_SEED in range(nrepeat) # 随机种子
     18 ]
     19 for temp_resi in tqdm(temp_reses, desc="Pipeline: "):
     20     ( # 函数的return
     21         mcmc_betas_mean, 
     22         mcmc_betas_std, 
     23         ols_betas_mean, 
     24         ols_betas_std
---> 25     ) = temp_resi.get()
     26     mcmc_betas_mean_l.append(mcmc_betas_mean)
     27     mcmc_betas_std_l.append(mcmc_betas_std)

File ~/miniforge3/envs/pymc_env0/lib/python3.11/multiprocessing/pool.py:774, in ApplyResult.get(self, timeout)
    772     return self._value
    773 else:
--> 774     raise self._value

Timeout: The file lock 'The file lock '/home/user/.pytensor/compiledir_Linux-3.10-el7.x86_64-x86_64-with-glibc2.17-x86_64-3.11.7-64/.lock' could not be acquired.
Apply node that caused the error: Composite{(i5 - (i2 + (i3 * i4) + (i0 * i1)))}(ExpandDims{axis=0}.0, [-4.045692 ... 74900e+00], ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, [0. 1. 1. ... 1. 0. 0.], Y_obs1{[-2.025834 ... 00053e+00]})
Toposort index: 6
Inputs types: [TensorType(float64, shape=(1,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1000,)), TensorType(float64, shape=(1000,))]

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.' could not be acquired.

Sorry my suggestion was unclear, the command is a context manager, you will want an with pytensor.config.... : above your call to pm.sample

Hi I tried to add this line in the beginning of a single simulation but got the error “No module named ‘numba’” Is it a typo or should I install any package?

def single_round_simulation(true_beta, RANDOM_SEED, samples_size, chain_len):
    with pytensor.config.change_flags(mode="NUMBA"):
        rng = np.random.default_rng(RANDOM_SEED)
        # skip some code because of it is too long
        basic_model = pm.Model()
        with basic_model:
            # similar to the user guide code
        with basic_model:
            idata = pm.sample(draws=chain_len, chains=1)
        return res
---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
                    ^^^^^^^^^^^^^^^^^^^
  File "/tmp/ipykernel_209230/1028081755.py", line 60, in single_round_simulation
    idata = pm.sample(draws=chain_len, chains=1)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 689, in sample
    step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 239, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 140, in instantiate_steppers
    step = step_class(vars=vars, model=model, **args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/hmc/nuts.py", line 180, in __init__
    super().__init__(vars, **kwargs)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/hmc/base_hmc.py", line 109, in __init__
    super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **pytensor_kwargs)
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/step_methods/arraystep.py", line 164, in __init__
    func = model.logp_dlogp_function(vars, dtype=dtype, **pytensor_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/model/core.py", line 618, in logp_dlogp_function
    return ValueGradFunction(costs, grad_vars, extra_vars_and_values, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/model/core.py", line 350, in __init__
    self._pytensor_function = compile_pymc(inputs, outputs, givens=givens, **kwargs)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pymc/pytensorf.py", line 991, in compile_pymc
    pytensor_function = pytensor.function(
                        ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/__init__.py", line 315, in function
    fn = pfunc(
         ^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py", line 469, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1762, in orig_function
    fn = m.create(defaults)
         ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1654, in create
    _fn, _i, _o = self.linker.make_thunk(
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/basic.py", line 245, in make_thunk
    return self.make_all(
           ^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/basic.py", line 688, in make_all
    thunks, nodes, jit_fn = self.create_jitable_thunk(
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/basic.py", line 638, in create_jitable_thunk
    converted_fgraph = self.fgraph_convert(
                       ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/numba/linker.py", line 25, in fgraph_convert
    from pytensor.link.numba.dispatch import numba_funcify
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/numba/dispatch/__init__.py", line 2, in <module>
    from pytensor.link.numba.dispatch.basic import numba_funcify, numba_typify
  File "/home/user/miniforge3/envs/pymc_env0/lib/python3.11/site-packages/pytensor/link/numba/dispatch/basic.py", line 10, in <module>
    import numba
ModuleNotFoundError: No module named 'numba'
"""

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

ModuleNotFoundError                       Traceback (most recent call last)
Cell In[15], line 24
     11 temp_reses = [ # return 
     12     pool.apply_async(
     13         single_round_simulation, # function
   (...)
     16     for RANDOM_SEED in range(nrepeat) # 随机种子
     17 ]
     18 for temp_resi in tqdm(temp_reses, desc="Pipeline: "):
     19     ( # 函数的return
     20         mcmc_betas_mean, 
     21         mcmc_betas_std, 
     22         ols_betas_mean, 
     23         ols_betas_std
---> 24     ) = temp_resi.get()
     25     mcmc_betas_mean_l.append(mcmc_betas_mean)
     26     mcmc_betas_std_l.append(mcmc_betas_std)

File ~/miniforge3/envs/pymc_env0/lib/python3.11/multiprocessing/pool.py:774, in ApplyResult.get(self, timeout)
    772     return self._value
    773 else:
--> 774     raise self._value

ModuleNotFoundError: No module named 'numba'

Yes you have to install numba, which is an optional package otherwise