Fix for issue # 813 modified the way kwargs are handled

[tanishy7777]

Solves issue #813

• Make sure all test pass.
• Make sure your code passes black.
• Make sure your code passes pylint.

[tomicapretto]

Hi @tanishy7777 thanks for working on this, and sorry for the delay in the response. I think we should do one of the following

• Emit a warning saying that some kwargs are not used by the JAX based sampler (as @GStechschulte mentions in the issue)
• Do nothing, which is the current behavior.

I'm not sure if the proposed changes are what we want. If I understand correctly, it will set up kwargs["num_steps"] to draws when it's not None, but would leave it unset when it's None (using whatever default value applies for that sampler in the library that implements it).

What if we have the following example:

kwargs = {
        "num_draws": 500,
}

blackjax_nuts_idata = model.fit(draws=250, inference_method="blackjax_nuts", **kwargs)

The user is passing through kwargs["num_draws"] they want 500 draws but we then replace that with 250, which is the one passed to draws=250.

This example makes me think we should just emit a warning in that case, and respect the name required by the underlying sampler.

---

A separate comment, it's always good to add a test when one implements a change in behavior like this one. Just let me/us know if you would need help with that.

[GStechschulte]

Thanks @tomicapretto for catching this.

bayeux uses the argument names (often num_chains for example) from the underlying libraries (blackjax, numpyro, tfp, etc.). However, if a library does not use that argument name, bayeux will silently ignore unused kwargs. Therefore, we may need to add logic according to the library to ensure the correct names are passed down.

Below is a comment and code snippet from Colin

The only samplers are from blackjax, numpyro, nutpie, tfp, and flowMC. bayeux uses the argument names from underlying libraries

for method in model.mcmc.methods:
  sampler = getattr(model.mcmc, method)
  print(method)
  for v in ('chains', 'num_chains', 'num_draws', 'num_samples', 'num_results'):
    if any(k == v for step in sampler.get_kwargs().values() for k in step):
      print('\t', v)

tfp_hmc
	 num_chains
	 num_draws
tfp_nuts
	 num_chains
	 num_draws
tfp_snaper_hmc
	 num_chains
	 num_results
blackjax_hmc
	 num_chains
	 num_draws
blackjax_chees_hmc
	 num_chains
	 num_draws
blackjax_meads_hmc
	 num_chains
	 num_draws
blackjax_nuts
	 num_chains
	 num_draws
blackjax_hmc_pathfinder
	 num_chains
	 num_draws
blackjax_nuts_pathfinder
	 num_chains
	 num_draws
numpyro_hmc
	 num_chains
	 num_samples
numpyro_nuts
	 num_chains
	 num_samples

As an aside, we should make sure changes here are compatible with #855. I do not think there is any major conflict, but good to double check.

[GStechschulte]

Hey @tanishy7777 any updates regarding my comment above? Thanks!

[tanishy7777]

Hey @tanishy7777 any updates regarding my comment above? Thanks!

Sorry for the late response, I think I might have missed this. Will look into it!

[tanishy7777]

Hey @tanishy7777 any updates regarding my comment above? Thanks!

If I understand correctly, you are saying that we should handle kwargs on a case by case basis, based on which library is being used(that we get from the inference_method parameter).

Fix 1:

So, if some argument is passed in kwargs which is not used in the underlying library we emit a warning to the user?

Fix 2a:

Also that we should map the extra arguments like draws -> num_draws based on the underlying library?

Fix 2b:

And give a warning to the user when both kwargs and an extra argument is passed (as for draws in this case), saying that num_draws in kwargs is being overridden by the arguments(draws) passed by them?

As show in #863 (comment)
where

kwargs = {
        "num_draws": 500,
}

blackjax_nuts_idata = model.fit(draws=250, inference_method="blackjax_nuts", **kwargs)

should produce warning because we use draws which overrides the num_draws from kwargs for blackjax_nuts