Weird VScode hangs on pymc v5 import on restart but not fresh install

Small update here.

It seems I am able to get pymc to successfully import by changing a pytensor environment variable as follows within my active conda environment:

conda env config vars set PYTENSOR_FLAGS="blas__ldflags=-framework Accelerate"

This is great as now my test script will print out the pymc version I am using. However, it not gets stuck on sampling with pm.sample(). It’s hard for me to know if this is a hang because of another unset or weirdly set environment variable but here is the following traceback:

(pymc) uqamcka3@x86_64-apple-darwin13 Random % /opt/miniconda3/envs/pymc/bin/p
ython /Users/uqamcka3/PHD/Projects/Random/scripts/multitest.py
*** Start script ***
*** Start pymc import ***
*** Start Model script ***
pymc: v. 5.16.2
^CTraceback (most recent call last):
  File "/Users/uqamcka3/PHD/Projects/Random/scripts/multitest.py", line 35, in <module>
    trace = pm.sample(chains=1, cores=1, random_seed=SEED)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 716, in sample
    step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 237, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 138, in instantiate_steppers
    step = step_class(vars=vars, model=model, **args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/step_methods/hmc/nuts.py", line 180, in __init__
    super().__init__(vars, **kwargs)
  File "/opt/miniconda3/envs/pymc/lib/python3.12/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 "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/step_methods/arraystep.py", line 163, in __init__
    func = model.logp_dlogp_function(vars, dtype=dtype, **pytensor_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/model/core.py", line 620, in logp_dlogp_function
    ip = self.initial_point(0)
         ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/model/core.py", line 1094, in initial_point
    fn = make_initial_point_fn(model=self, return_transformed=True)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/initial_point.py", line 152, in make_initial_point_fn
    func = compile_pymc(inputs=[], outputs=initial_values, mode=pytensor.compile.mode.FAST_COMPILE)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pymc/pytensorf.py", line 1039, in compile_pymc
    pytensor_function = pytensor.function(
                        ^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/compile/function/__init__.py", line 318, in function
    fn = pfunc(
         ^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/compile/function/pfunc.py", line 465, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 1750, in orig_function
    m = Maker(
        ^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 1523, in __init__
    self.prepare_fgraph(inputs, outputs, found_updates, fgraph, mode, profile)
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 1411, in prepare_fgraph
    rewriter_profile = rewriter(fgraph)
                       ^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 125, in __call__
    return self.rewrite(fgraph)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 121, in rewrite
    return self.apply(fgraph, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 291, in apply
    sub_prof = rewriter.apply(fgraph)
               ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 2443, in apply
    sub_prof = grewrite.apply(fgraph)
               ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 2027, in apply
    nb += self.process_node(fgraph, node)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/tensor/rewriting/basic.py", line 1122, in constant_folding
    thunk = node.op.make_thunk(node, storage_map, compute_map, no_recycling=[])
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/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 "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/op.py", line 84, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/basic.py", line 1182, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/basic.py", line 1103, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/basic.py", line 1614, in cthunk_factory
    key = self.cmodule_key()
          ^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/basic.py", line 1266, in cmodule_key
    compile_args=self.compile_args(),
                 ^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/basic.py", line 952, in compile_args
    ret += c_compiler.compile_args()
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/cmodule.py", line 2315, in compile_args
    compilation_result, execution_result = try_march_flag(
                                           ^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/cmodule.py", line 2044, in try_march_flag
    compilation_result, execution_result = GCC_compiler.try_compile_tmp(
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/cmodule.py", line 2407, in try_compile_tmp
    return cls._try_compile_tmp(
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/link/c/cmodule.py", line 1911, in _try_compile_tmp
    out, err, p_ret = output_subprocess_Popen([exe_path])
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/site-packages/pytensor/utils.py", line 200, in output_subprocess_Popen
    out = p.communicate()
          ^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/subprocess.py", line 1209, in communicate
    stdout, stderr = self._communicate(input, endtime, timeout)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/subprocess.py", line 2115, in _communicate
    ready = selector.select(timeout)
            ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/pymc/lib/python3.12/selectors.py", line 415, in select
    fd_event_list = self._selector.poll(timeout)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt

I am unsure what to make of this, the main thing I can conclude is that it at least gets stuck before the initiation of the NUTS sampler. I have also tried importing pytensor and manually setting my clang path as in this issue, which didn’t seem to help.

this is the test script I am running:

#%%
print('*** Start script ***')
import numpy as np
import pandas as pd
import arviz as az
#%%
print('*** Start pymc import ***')
import pymc as pm

print('*** Start Model script ***')
print(f'{pm.__name__}: v. {pm.__version__}')

if __name__ == '__main__':
    SEED = 20180730
    np.random.seed(SEED)

    # Generate data
    mu_real = 0
    sd_real = 1
    n_samples = 1000
    y = np.random.normal(loc=mu_real, scale=sd_real, size=n_samples)

    # Bayesian modelling
    with pm.Model() as model:
    
        mu = pm.Normal('mu', mu=0, sigma=10)
        sd = pm.HalfNormal('sd', sigma=10)
    
        # Likelihood
        likelihood = pm.Normal('likelihood', mu=mu, sigma=sd, observed=y)    
        trace = pm.sample(chains=1, cores=1, random_seed=SEED)

    print('Done!')

    az.plot_trace(trace, figsize=(10,7)) 
# %%

Hopefully, this might be a bit more helpful for troubleshooting. Let me know if there is anything else I can provide to make troubleshooting easier