code: def test_save_warmup(self, save_warmup, chains, tune, draws):
      with pm.Model():
          pm.Uniform("u1")
          pm.Normal("n1")
          with warnings.catch_warnings():
              warnings.filterwarnings(
                  "ignore", ".*number of samples.*", UserWarning
              )
              warnings.filterwarnings(
                  "ignore", "More chains .* than draws.*", UserWarning
              )
              idata = pm.sample(
                  tune=tune,
                  draws=draws,
                  chains=chains,
                  cores=1,
                  step=pm.Metropolis(),
                  discard_tuned_samples=False,
                  return_inferencedata=True,
                  idata_kwargs={"save_warmup": save_warmup},
              )
      warmup_prefix = "" if save_warmup and (tune > 0) else "~"
      post_prefix = "" if draws > 0 else "~"
      test_dict = {
          f"{post_prefix}posterior": ["u1", "n1"],
          f"{post_prefix}sample_stats": ["~tune", "accept"],
          f"{warmup_prefix}warmup_posterior": ["u1", "n1"],
          f"{warmup_prefix}warmup_sample_stats": ["~tune"],
          "~warmup_log_likelihood": [],
          "~log_likelihood": [],
      }
      fails = check_multiple_attrs(test_dict, idata)
      assert not fails
      if hasattr(idata, "posterior"):
          assert idata.posterior.sizes["chain"] == chains
          assert idata.posterior.sizes["draw"] == draws
      if hasattr(idata, "warmup_posterior"):
          assert idata.warmup_posterior.sizes["chain"] == chains
          assert idata.warmup_posterior.sizes["draw"] == tune
  
imports: import warnings
  from typing import Dict, Tuple
  import numpy
  import pytensor.tensor
  import pytest
  from arviz import InferenceData
  from arviz.tests.helpers import check_multiple_attrs
  from numpy import ma
  from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
  import pymc
  from pymc.backends.arviz import InferenceDataConverter, predictions_to_inference_data, to_inference_data
  from pymc.exceptions import ImputationWarning
label_desc: Use image node_id 1 for calling the TestPyMCWarmupHandling obj's underlying member method code with example usage: obj.test_save_warmup(save_warmup, chains, tune, draws) without return types