Sampler argument target_accept not working in 4.0.0b4

[larryshamalama]

with pm.Model() as model:
    norm = pm.Normal("norm")
    pm.sample(target_accept=0.9, tune=2000)

yields

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [2], in <module>
      1 with pm.Model() as model:
      2     norm = pm.Normal("norm")
----> 3     pm.sample(target_accept=0.9, tune=2000)

File ~/opt/anaconda3/envs/pymc-dev-py39/lib/python3.9/site-packages/pymc/sampling.py:470, in sample(draws, step, init, n_init, initvals, trace, chain_idx, chains, cores, tune, progressbar, model, random_seed, discard_tuned_samples, compute_convergence_checks, callback, jitter_max_retries, return_inferencedata, idata_kwargs, mp_ctx, **kwargs)
    467 draws += tune
    469 initial_points = None
--> 470 step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
    472 if isinstance(step, list):
    473     step = CompoundStep(step)

File ~/opt/anaconda3/envs/pymc-dev-py39/lib/python3.9/site-packages/pymc/sampling.py:219, in assign_step_methods(model, step, methods, step_kwargs)
    211         selected = max(
    212             methods,
    213             key=lambda method, var=rv_var, has_gradient=has_gradient: method._competence(
    214                 var, has_gradient
    215             ),
    216         )
    217         selected_steps[selected].append(var)
--> 219 return instantiate_steppers(model, steps, selected_steps, step_kwargs)

File ~/opt/anaconda3/envs/pymc-dev-py39/lib/python3.9/site-packages/pymc/sampling.py:142, in instantiate_steppers(model, steps, selected_steps, step_kwargs)
    140 unused_args = set(step_kwargs).difference(used_keys)
    141 if unused_args:
--> 142     raise ValueError("Unused step method arguments: %s" % unused_args)
    144 if len(steps) == 1:
    145     return steps[0]

ValueError: Unused step method arguments: {'target_accept'}

with PyMC version 4.0.0b4, aesara version 2.5.1 and aeppl version 0.0.27. However, there is no error when using PyMC version 4.0.0b3 (and aesara version 2.4.0, aeppl version 0.0.26). See notebooks here (b3) and here (b4)

First raised by @OriolAbril

[aloctavodia]

This was introduced here 2dd4c8c now all steps are assigned using assign_step_methods but before that we first attempted to run NUTS and only if that fails we used assign_step_methods.

Following this change parameters specific to a given sampler has to be passed using the name of the step method, like nuts={'target_accept:0.9'}

One solution is to keep the previous change and just update the documentation.

[ricardoV94]

Maybe we can try and assume the kwargs are meant for nuts, since that's the most common case?

Would need some extra logic, but perhaps nothing too convoluted?

[aloctavodia]

Yes, I think that is a good solution. I can work on it.