BUG: nuts vectorized chain method arg raising in pymc>=6

Describe the issue:

Since pymc 6.x, pm.sample no longer forwards chain_method (or any other top-level parameter of pymc.sampling.jax.sample_jax_nuts, e.g. postprocessing_backend) to the external JAX samplers. Every spelling is either misrouted or silently dropped:

• pm.sample(..., nuts_sampler="blackjax", nuts_sampler_kwargs={"chain_method": "vectorized"}) (the pymc 5.x form): the deprecation shim converts nuts_sampler_kwargs into nuts={...} (pymc/sampling/mcmc.py:847), and nuts kwargs are passed to the NUTS kernel via nuts_kwargs, so blackjax raises TypeError: build_kernel.<locals>.kernel() got an unexpected keyword argument 'chain_method' during warmup.
• pm.sample(..., nuts={"chain_method": "vectorized"}) (the documented replacement): same TypeError.
• pm.sample(..., chain_method="vectorized") (top-level): swallowed silently by _sample_external_nuts(**kwargs) (pymc/sampling/mcmc.py:372), which never forwards its **kwargs to sample_jax_nuts (pymc/sampling/mcmc.py:512). No warning, no error, chain_method stays at its default "parallel".

In pymc 5.x, _sample_external_nuts spread **nuts_sampler_kwargs as top-level arguments of sample_jax_nuts, which is how chain_method used to reach the sampler.

The consequence: with nuts_sampler="blackjax", vectorized chains are unreachable through pm.sample, and the default chain_method="parallel" crashes on the common single-device case whenever chains > len(jax.devices()) (pmap error, surfacing as ValueError: cannot select an axis to squeeze out which has size not equal to one). The same routing applies to nuts_sampler="numpyro" (its default degrades gracefully to sequential chains, but chain_method is equally unsettable).

Calling pymc.sampling.jax.sample_blackjax_nuts(chain_method="vectorized", ...) directly still works (with blackjax <= 1.5), which suggests the routing change in pm.sample is unintentional.

Reproduceable code example:

import numpy as np
import pymc as pm

with pm.Model():
    x = pm.Normal("x", 0.0, 1.0)
    pm.Normal("obs", x, 1.0, observed=np.array([0.3, -0.1, 0.5]))
    idata = pm.sample(
        draws=50, tune=50, chains=4,
        nuts_sampler="blackjax",
        nuts={"chain_method": "vectorized"},
        progressbar=False,
    )

Error message:

TypeError: build_kernel.<locals>.kernel() got an unexpected keyword argument 'chain_method'


(from inside `blackjax/adaptation/window_adaptation.py` during warmup; with the top-level spelling there is no error at all, `chain_method` is silently ignored)

PyMC version information:

pymc 6.1.0, pytensor (bundled), blackjax 1.5 / 1.6 (both affected), jax 0.10.x, Python 3.12, macOS (arm64) and Linux. Works as documented on pymc 5.27.1 + blackjax 1.5.

Context for the issue:

No response

Start in pymc/sampling/mcmc.py at the deprecation shim around line 847 and _sample_external_nuts, especially the call near line 512. Trace how nuts, nuts_sampler_kwargs, and top-level arguments reach sample_jax_nuts, then run the provided four-chain BlackJAX reproducer. Done means chain_method="vectorized" reaches the external sampler for both BlackJAX and NumPyro without being passed to the NUTS kernel or silently ignored.