code: def init_nuts(
      *,
      init: str = "auto",
      chains: int = 1,
      n_init: int = 500_000,
      model: Optional[Model] = None,
      random_seed: RandomSeed = None,
      progressbar=True,
      jitter_max_retries: int = 10,
      tune: Optional[int] = None,
      initvals: Optional[
          Union[StartDict, Sequence[Optional[StartDict]]]
      ] = None,
      **kwargs,
  ) -> Tuple[Sequence[PointType], NUTS]:
      """Set up the mass matrix initialization for NUTS.
  
      NUTS convergence and sampling speed is extremely dependent on the
      choice of mass/scaling matrix. This function implements different
      methods for choosing or adapting the mass matrix.
  
      Parameters
      ----------
      init : str
          Initialization method to use.
  
          * auto: Choose a default initialization method automatically.
            Currently, this is ``jitter+adapt_diag``, but this can change in the future. If you
            depend on the exact behaviour, choose an initialization method explicitly.
          * adapt_diag: Start with a identity mass matrix and then adapt a diagonal based on the
            variance of the tuning samples. All chains use the test value (usually the prior mean)
            as starting point.
          * jitter+adapt_diag: Same as ``adapt_diag``, but use test value plus a uniform jitter in
            [-1, 1] as starting point in each chain.
          * jitter+adapt_diag_grad:
            An experimental initialization method that uses information from gradients and samples
            during tuning.
          * advi+adapt_diag: Run ADVI and then adapt the resulting diagonal mass matrix based on the
            sample variance of the tuning samples.
          * advi: Run ADVI to estimate posterior mean and diagonal mass matrix.
          * advi_map: Initialize ADVI with MAP and use MAP as starting point.
          * map: Use the MAP as starting point. This is discouraged.
          * adapt_full: Adapt a dense mass matrix using the sample covariances. All chains use the
            test value (usually the prior mean) as starting point.
          * jitter+adapt_full: Same as ``adapt_full``, but use test value plus a uniform jitter in
            [-1, 1] as starting point in each chain.
  
      chains : int
          Number of jobs to start.
      initvals : optional, dict or list of dicts
          Dict or list of dicts with initial values to use instead of the defaults from `Model.initial_values`.
          The keys should be names of transformed random variables.
      n_init : int
          Number of iterations of initializer. Only works for 'ADVI' init methods.
      model : Model (optional if in ``with`` context)
      random_seed : int, array-like of int, RandomState or Generator, optional
          Seed for the random number generator.
      progressbar : bool
          Whether or not to display a progressbar for advi sampling.
      jitter_max_retries : int
          Maximum number of repeated attempts (per chain) at creating an initial matrix with uniform jitter
          that yields a finite probability. This applies to ``jitter+adapt_diag`` and ``jitter+adapt_full``
          init methods.
      **kwargs : keyword arguments
          Extra keyword arguments are forwarded to pymc.NUTS.
  
      Returns
      -------
      initial_points : list
          Starting points for each chain.
      nuts_sampler : ``pymc.step_methods.NUTS``
          Instantiated and initialized NUTS sampler object
      """
      model = modelcontext(model)
  
      vars = kwargs.get("vars", model.value_vars)
      if set(vars) != set(model.value_vars):
          raise ValueError(
              "Must use init_nuts on all variables of a model."
          )
      if not all_continuous(vars):
          raise ValueError(
              "init_nuts can only be used for models with continuous variables."
          )
  
      if not isinstance(init, str):
          raise TypeError("init must be a string.")
  
      init = init.lower()
  
      if init == "auto":
          init = "jitter+adapt_diag"
  
      random_seed_list = _get_seeds_per_chain(random_seed, chains)
  
      _log.info(f"Initializing NUTS using {init}...")
  
      cb = [
          pm.callbacks.CheckParametersConvergence(
              tolerance=1e-2, diff="absolute"
          ),
          pm.callbacks.CheckParametersConvergence(
              tolerance=1e-2, diff="relative"
          ),
      ]
  
      initial_points = _init_jitter(
          model,
          initvals,
          seeds=random_seed_list,
          jitter="jitter" in init,
          jitter_max_retries=jitter_max_retries,
      )
  
      apoints = [
          DictToArrayBijection.map(point) for point in initial_points
      ]
      apoints_data = [apoint.data for apoint in apoints]
      potential: quadpotential.QuadPotential
  
      if init == "adapt_diag":
          mean = np.mean(apoints_data, axis=0)
          var = np.ones_like(mean)
          n = len(var)
          potential = quadpotential.QuadPotentialDiagAdapt(
              n, mean, var, 10
          )
      elif init == "jitter+adapt_diag":
          mean = np.mean(apoints_data, axis=0)
          var = np.ones_like(mean)
          n = len(var)
          potential = quadpotential.QuadPotentialDiagAdapt(
              n, mean, var, 10
          )
      elif init == "jitter+adapt_diag_grad":
          mean = np.mean(apoints_data, axis=0)
          var = np.ones_like(mean)
          n = len(var)
  
          if tune is not None and tune > 250:
              stop_adaptation = tune - 50
          else:
              stop_adaptation = None
  
          potential = quadpotential.QuadPotentialDiagAdaptExp(
              n,
              mean,
              alpha=0.02,
              use_grads=True,
              stop_adaptation=stop_adaptation,
          )
      elif init == "advi+adapt_diag":
          approx = pm.fit(
              random_seed=random_seed_list[0],
              n=n_init,
              method="advi",
              model=model,
              callbacks=cb,
              progressbar=progressbar,
              obj_optimizer=pm.adagrad_window,
          )
          approx_sample = approx.sample(
              draws=chains,
              random_seed=random_seed_list[0],
              return_inferencedata=False,
          )
          initial_points = [approx_sample[i] for i in range(chains)]
          std_apoint = approx.std.eval()
          cov = std_apoint**2
          mean = approx.mean.get_value()
          weight = 50
          n = len(cov)
          potential = quadpotential.QuadPotentialDiagAdapt(
              n, mean, cov, weight
          )
      elif init == "advi":
          approx = pm.fit(
              random_seed=random_seed_list[0],
              n=n_init,
              method="advi",
              model=model,
              callbacks=cb,
              progressbar=progressbar,
              obj_optimizer=pm.adagrad_window,
          )
          approx_sample = approx.sample(
              draws=chains,
              random_seed=random_seed_list[0],
              return_inferencedata=False,
          )
          initial_points = [approx_sample[i] for i in range(chains)]
          cov = approx.std.eval() ** 2
          potential = quadpotential.QuadPotentialDiag(cov)
      elif init == "advi_map":
          start = pm.find_MAP(
              include_transformed=True, seed=random_seed_list[0]
          )
          approx = pm.MeanField(model=model, start=start)
          pm.fit(
              random_seed=random_seed_list[0],
              n=n_init,
              method=pm.KLqp(approx),
              callbacks=cb,
              progressbar=progressbar,
              obj_optimizer=pm.adagrad_window,
          )
          approx_sample = approx.sample(
              draws=chains,
              random_seed=random_seed_list[0],
              return_inferencedata=False,
          )
          initial_points = [approx_sample[i] for i in range(chains)]
          cov = approx.std.eval() ** 2
          potential = quadpotential.QuadPotentialDiag(cov)
      elif init == "map":
          start = pm.find_MAP(
              include_transformed=True, seed=random_seed_list[0]
          )
          cov = pm.find_hessian(point=start)
          initial_points = [start] * chains
          potential = quadpotential.QuadPotentialFull(cov)
      elif init == "adapt_full":
          mean = np.mean(apoints_data * chains, axis=0)
          initial_point = initial_points[0]
          initial_point_model_size = sum(
              initial_point[n.name].size for n in model.value_vars
          )
          cov = np.eye(initial_point_model_size)
          potential = quadpotential.QuadPotentialFullAdapt(
              initial_point_model_size, mean, cov, 10
          )
      elif init == "jitter+adapt_full":
          mean = np.mean(apoints_data, axis=0)
          initial_point = initial_points[0]
          initial_point_model_size = sum(
              initial_point[n.name].size for n in model.value_vars
          )
          cov = np.eye(initial_point_model_size)
          potential = quadpotential.QuadPotentialFullAdapt(
              initial_point_model_size, mean, cov, 10
          )
      else:
          raise ValueError(f"Unknown initializer: {init}.")
  
      step = pm.NUTS(potential=potential, model=model, **kwargs)
  
      # Filter deterministics from initial_points
      value_var_names = [var.name for var in model.value_vars]
      initial_points = [
          {
              k: v
              for k, v in initial_point.items()
              if k in value_var_names
          }
          for initial_point in initial_points
      ]
  
      return initial_points, step
  
imports: import logging
  import pickle
  import sys
  import time
  import warnings
  from typing import Any, Dict, Iterator, List, Literal, Mapping, Optional, Sequence, Set, Tuple, Type, Union, overload
  import numpy
  import pytensor.gradient
  from arviz import InferenceData, dict_to_dataset
  from arviz.data.base import make_attrs
  from fastprogress.fastprogress import progress_bar
  from pytensor.graph.basic import Variable
  from typing_extensions import Protocol, TypeAlias
  import pymc
  from pymc.backends import RunType, TraceOrBackend, init_traces
  from pymc.backends.arviz import coords_and_dims_for_inferencedata, find_constants, find_observations
  from pymc.backends.base import IBaseTrace, MultiTrace, _choose_chains
  from pymc.blocking import DictToArrayBijection
  from pymc.exceptions import SamplingError
  from pymc.initial_point import PointType, StartDict, make_initial_point_fns_per_chain
  from pymc.model import Model, modelcontext
  from pymc.sampling.parallel import Draw, _cpu_count
  from pymc.sampling.population import _sample_population
  from pymc.stats.convergence import log_warning_stats, log_warnings, run_convergence_checks
  from pymc.step_methods import NUTS, CompoundStep
  from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
  from pymc.step_methods.hmc import quadpotential
  from pymc.util import RandomSeed, RandomState, _get_seeds_per_chain, drop_warning_stat, get_untransformed_name, is_transformed_name
  from pymc.vartypes import discrete_types
label_desc: Use image node_id 16 for calling a global function with example usage: init_nuts() and returns: initial_points,  step