Fix error when passing coords and dims in sampling_jax

[bherwerth]

Fix #5932 : Allow for passing coords and dims within idata_kwargs in sample_blackjax_nuts and sample_numpyro_nuts.

Previously, the arguments dims and coords as defined in the sampling function were passed to az.from_dict along with idata_kwargs. This caused an error when idata_kwargs contained keys coords and/or dims (duplicate keyword arguments), cf #5932.

Proposed solution:

• By using functools.partial on az.from_dict, new defaults are set first. Only afterwards idata_kwargs is passed along. This results in dims and coords in idata_kwargs taking precedence (if given) instead of causing duplicate keyword args.
• Test cases were added to test_idata_kwargs for providing coords and/or dims in idata_kwargs.
  ...

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• ...

Bugfixes / New features

• Fix Cannot pass dims to idata_kwargs parameter in sample_numpyro_nuts #5932

Docs / Maintenance

• ...

[codecov]

Codecov Report

Merging #5983 (13e6ed1) into main (2583b7f) will decrease coverage by 0.00%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #5983      +/-   ##
==========================================
- Coverage   89.35%   89.35%   -0.01%     
==========================================
  Files          73       73              
  Lines       13251    13253       +2     
==========================================
+ Hits        11841    11842       +1     
- Misses       1410     1411       +1     

Impacted Files	Coverage Δ
pymc/sampling_jax.py	96.96% <100.00%> (+0.03%)	⬆️
pymc/step_methods/hmc/base_hmc.py	89.76% <0.00%> (-0.79%)	⬇️