BUG: keep_warning_stat argument in pm.sample does nothing

[ivlis]

Describe the issue:

For nutpie sample divergence_message stays in the idata object regardless of the setting what breaks to_netcdf()

https://www.pymc.io/projects/docs/en/stable/api/generated/pymc.sample.html

Reproduceable code example:

with pm.Model() as m:
    a = pm.Normal('a')
    idata = pm.sample(keep_warning_stat=False, nuts_sampler="nutpie")

print(f"nutpie, False: divergence_message in idata: {'divergence_message' in idata['sample_stats']}")

with pm.Model() as m:
    a = pm.Normal('a')
    idata = pm.sample(keep_warning_stat=True, nuts_sampler="nutpie")

print(f"nutpie, True: divergence_message in idata: {'divergence_message' in idata['sample_stats']}")


with pm.Model() as m:
    a = pm.Normal('a')
    idata = pm.sample(keep_warning_stat=False, nuts_sampler="pymc")

print(f"pymc, False: divergence_message in idata: {'divergence_message' in idata['sample_stats']}")

with pm.Model() as m:
    a = pm.Normal('a')
    idata = pm.sample(keep_warning_stat=True, nuts_sampler="pymc")

print(f"pymc, True: divergence_message in idata: {'divergence_message' in idata['sample_stats']}")


Output

NUTS[nutpie]: [a]

nutpie, False: divergence_message in idata: True

NUTS[nutpie]: [a]

nutpie, True: divergence_message in idata: True

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [a]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 0 seconds.

pymc, False: divergence_message in idata: False

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [a]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 0 seconds.

pymc, True: divergence_message in idata: False

Error message:

PyMC version information:

6.0.1

Context for the issue:

No response

[ricardoV94]

Slipped through the cracks... so many arguments. We should fix

[ricardoV94]

I didn't even know we compute convergence checks when running with nutpie

[omesh-s]

Hi! I’d like to work on this as my first contribution (as part of a CodePath open‑source project). My plan is to mirror the keep_warning_stat behavior from the pymc sampler in the nutpie sampler path and add a regression test using the example in the issue description. Is it okay if I take this?