Error KeyError: 'diverging' when setting up CategoricalGibbsMetropolis step method with proposal='proportional'

I am trying to set up CategoricalGibbsMetropolis step method for a categorical variable with the proposal parameter set to proportional as follows:

    with model:

        # Create and save the model graph.
        graph = pm.model_to_graphviz(model)
        filename = f"Model_graphs/{model_name}_graph"
        graph.render(filename, format="pdf", cleanup=True)

        # Debug the model before sampling.
        model.debug()
        
        # Set up the categorical sampler for discrete variables.
        step1 = pm.CategoricalGibbsMetropolis(
            vars=[model.gp], 
            proposal='proportional'
        )

        # Set up the NUTS sampler for continuous variables.
        step2 = pm.NUTS(
            vars=[var for var in model.free_RVs if var.name != 'gp'],
            target_accept=0.99, 
            max_treedepth=15
        )

        # Compute trace (posterior).
        idata = pm.sample(
            draws=2000,
            tune=2500,
            chains=4,
            cores=4,
            step=[step1, step2],
            return_inferencedata=True,
            init="jitter+adapt_diag",
            compute_convergence_checks=True,
            discard_tuned_samples=True,
        )

I am encountering an error just as the sampling start only if i set proposal=β€˜proportional’ in the CategoricalGibbsMetropolis step method:

Multiprocess sampling (4 chains in 4 jobs)
CompoundStep
>CategoricalGibbsMetropolis: [gp]
>NUTS: [pi]
Python Interpreter Path: C:\Users\fbabb\.conda\envs\pymc_env\python.exe
Python Version: 3.12.7
PyMC version: 5.19.1
Python Interpreter Path: C:\Users\fbabb\.conda\envs\pymc_env\python.exe
Python Version: 3.12.7
PyMC version: 5.19.1
Python Interpreter Path: C:\Users\fbabb\.conda\envs\pymc_env\python.exe
Python Version: 3.12.7
PyMC version: 5.19.1
Python Interpreter Path: C:\Users\fbabb\.conda\envs\pymc_env\python.exe
Python Version: 3.12.7
PyMC version: 5.19.1
Sampling 4 chains, 0 divergences ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━   0% -:--:-- / 0:00:00
Traceback (most recent call last):
  File "h:\Fear_Generalization\run_pymc.py", line 154, in <module>
    Result_Study1_CLG2D = sampling_fun(
                          ^^^^^^^^^^^^^
  File "h:\Fear_Generalization\run_pymc.py", line 122, in sampling_fun
    idata = pm.sample(
            ^^^^^^^^^^
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py", line 906, in sample
    _mp_sample(**sample_args, **parallel_args)
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py", line 1318, in _mp_sample
    strace.record(draw.point, draw.stats)
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\backends\ndarray.py", line 116, in record
    data[key][draw_idx] = val
    ~~~~^^^^^
KeyError: 'diverging'

Starting sampling while leaving proposal = β€˜uniform’ does not cause any error.
I was wondering if the proportional proposal is still supported, and if yes if i am doing something wrong.
My current version of PyMC is 5.19.1 but the same error occured with version 5.18.2 thats why i tried to resolve the issue by updating.

Looks like a bug, can you provide a more minimal fully reproducible example?

1 Like

From a quick look at the source code, it seems like CategoricalGibbsMetropolis.astep_prop is not returning the stats as is done in astep_unif. A reproducible example would really help though, because I don’t understand why the β€œdiverging” stat from NUTS is the one that leads to trouble.

1 Like

@mercer, maybe you could try editing your pymc source by copying this pattern at the end of astep_prop and see if that fixes your problem. If it does, please open an issue on github with a minimal model that causes the problem, and we’ll apply the patch (or you can open a PR if you’re feel up to it).

1 Like

Thank you, that resolved the error for me. I opened the issue on github.