BUG: KeyError: 'diverging' when setting up CategoricalGibbsMetropolis step method with proposal=‘proportional’

[MarcelloCalza]

Describe the issue:

Trying to set up CategoricalGibbsMetropolis step method for a categorical variable with the proposal parameter set to 'proportional' causes an error KeyError: 'diverging' just as sampling starts.

Reproduceable code example:

import pymc as pm

with pm.Model() as model:
    # Prior for categorical
    p = pm.Dirichlet('p', a=[1., 1., 1.], shape=3)

    # Categorical variable
    category = pm.Categorical('category', p=p, shape=100)

    # CategoricalGibbsMetropolis with proportional proposal
    step = pm.CategoricalGibbsMetropolis(vars=[category], proposal = 'proportional')

    trace = pm.sample(500, cores=1, chains=1, step=step, return_inferencedata=False)

Error message:

Traceback (most recent call last):
  File "h:\Fear_Generalization\import pymc as pm.py", line 14, in <module>
    trace = pm.sample(500, cores=1, chains=1, step=step, return_inferencedata=False)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py", line 929, in sample
    _sample_many(**sample_args)
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py", line 1072, in _sample_many
    _sample(
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py", line 1157, in _sample
    for it, diverging in enumerate(sampling_gen):
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\fbabb\.conda\envs\pymc_env\Lib\site-packages\pymc\sampling\mcmc.py", line 1224, in _iter_sample
    trace.record(point, 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'

PyMC version information:

PyMC Version: 5.19.1
PyTensor Version: 2.26.3
Python Version: 3.12.7

Context for the issue:

Editing the pymc source code by copying this pattern at the end of astep_prop as suggested by Luciano Paz resolved the issue for me.

[welcome]

]
🎉 Welcome to PyMC! 🎉 We're really excited to have your input into the project! 💖

If you haven't done so already, please make sure you check out our Contributing Guidelines and Code of Conduct.