original_docstring: 
          Run the MCMC samplers and collect samples.
  
          :param random.PRNGKey rng_key: Random number generator key to be used for the sampling.
              For multi-chains, a batch of `num_chains` keys can be supplied. If `rng_key`
              does not have batch_size, it will be split in to a batch of `num_chains` keys.
          :param args: Arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init` method.
              These are typically the arguments needed by the `model`.
          :param extra_fields: Extra fields (aside from `"z"`, `"diverging"`) to be collected
              during the MCMC run. Note that subfields can be accessed using dots, e.g.
              `"adapt_state.step_size"` can be used to collect step sizes at each step.
          :type extra_fields: tuple or list of str
          :param init_params: Initial parameters to begin sampling. The type must be consistent
              with the input type to `potential_fn`.
          :param kwargs: Keyword arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init`
              method. These are typically the keyword arguments needed by the `model`.
  
          .. note:: jax allows python code to continue even when the compiled code has not finished yet.
              This can cause troubles when trying to profile the code for speed.
              See https://jax.readthedocs.io/en/latest/async_dispatch.html and
              https://jax.readthedocs.io/en/latest/profiling.html for pointers on profiling jax programs.
          
docstring: Run the MCMC samplers and collect samples.
docstring_tokens: Run the MCMC samplers and collect samples .
code: def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
          """
          Run the MCMC samplers and collect samples.
  
          :param random.PRNGKey rng_key: Random number generator key to be used for the sampling.
              For multi-chains, a batch of `num_chains` keys can be supplied. If `rng_key`
              does not have batch_size, it will be split in to a batch of `num_chains` keys.
          :param args: Arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init` method.
              These are typically the arguments needed by the `model`.
          :param extra_fields: Extra fields (aside from `"z"`, `"diverging"`) to be collected
              during the MCMC run. Note that subfields can be accessed using dots, e.g.
              `"adapt_state.step_size"` can be used to collect step sizes at each step.
          :type extra_fields: tuple or list of str
          :param init_params: Initial parameters to begin sampling. The type must be consistent
              with the input type to `potential_fn`.
          :param kwargs: Keyword arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init`
              method. These are typically the keyword arguments needed by the `model`.
  
          .. note:: jax allows python code to continue even when the compiled code has not finished yet.
              This can cause troubles when trying to profile the code for speed.
              See https://jax.readthedocs.io/en/latest/async_dispatch.html and
              https://jax.readthedocs.io/en/latest/profiling.html for pointers on profiling jax programs.
          """
          init_params = tree_map(
              lambda x: lax.convert_element_type(x, jnp.result_type(x)), init_params
          )
          self._args = args
          self._kwargs = kwargs
          init_state = self._get_cached_init_state(rng_key, args, kwargs)
          if self.num_chains > 1 and rng_key.ndim == 1:
              rng_key = random.split(rng_key, self.num_chains)
  
          if self._warmup_state is not None:
              self._set_collection_params(0, self.num_samples, self.num_samples, "sample")
              init_state = self._warmup_state._replace(rng_key=rng_key)
  
          if init_params is not None and self.num_chains > 1:
              prototype_init_val = tree_flatten(init_params)[0][0]
              if jnp.shape(prototype_init_val)[0] != self.num_chains:
                  raise ValueError(
                      "`init_params` must have the same leading dimension"
                      " as `num_chains`."
                  )
          assert isinstance(extra_fields, (tuple, list))
          collect_fields = tuple(
              set(
                  (self._sample_field,)
                  + tuple(self._default_fields)
                  + tuple(extra_fields)
              )
          )
          partial_map_fn = partial(
              self._single_chain_mcmc,
              args=args,
              kwargs=kwargs,
              collect_fields=collect_fields,
          )
          map_args = (rng_key, init_state, init_params)
          if self.num_chains == 1:
              states_flat, last_state = partial_map_fn(map_args)
              states = tree_map(lambda x: x[jnp.newaxis, ...], states_flat)
          else:
              if self.chain_method == "sequential":
                  states, last_state = _laxmap(partial_map_fn, map_args)
              elif self.chain_method == "parallel":
                  states, last_state = pmap(partial_map_fn)(map_args)
              else:
                  assert self.chain_method == "vectorized"
                  states, last_state = partial_map_fn(map_args)
                  # swap num_samples x num_chains to num_chains x num_samples
                  states = tree_map(lambda x: jnp.swapaxes(x, 0, 1), states)
              states_flat = tree_map(
                  lambda x: jnp.reshape(x, (-1,) + x.shape[2:]), states
              )
          self._last_state = last_state
          self._states = states
          self._states_flat = states_flat
          self._set_collection_params()
code_tokens: def run ( self , rng_key , * args , extra_fields = ( ) , init_params = None , ** kwargs ) : init_params = tree_map ( lambda x : lax . convert_element_type ( x , jnp . result_type ( x ) ) , init_params ) self . _args = args self . _kwargs = kwargs init_state = self . _get_cached_init_state ( rng_key , args , kwargs ) if self . num_chains > 1 and rng_key . ndim == 1 : rng_key = random . split ( rng_key , self . num_chains ) if self . _warmup_state is not None : self . _set_collection_params ( 0 , self . num_samples , self . num_samples , "sample" ) init_state = self . _warmup_state . _replace ( rng_key = rng_key ) if init_params is not None and self . num_chains > 1 : prototype_init_val = tree_flatten ( init_params ) [ 0 ] [ 0 ] if jnp . shape ( prototype_init_val ) [ 0 ] != self . num_chains : raise ValueError ( "`init_params` must have the same leading dimension" " as `num_chains`." ) assert isinstance ( extra_fields , ( tuple , list ) ) collect_fields = tuple ( set ( ( self . _sample_field , ) + tuple ( self . _default_fields ) + tuple ( extra_fields ) ) ) partial_map_fn = partial ( self . _single_chain_mcmc , args = args , kwargs = kwargs , collect_fields = collect_fields , ) map_args = ( rng_key , init_state , init_params ) if self . num_chains == 1 : states_flat , last_state = partial_map_fn ( map_args ) states = tree_map ( lambda x : x [ jnp . newaxis , ... ] , states_flat ) else : if self . chain_method == "sequential" : states , last_state = _laxmap ( partial_map_fn , map_args ) elif self . chain_method == "parallel" : states , last_state = pmap ( partial_map_fn ) ( map_args ) else : assert self . chain_method == "vectorized" states , last_state = partial_map_fn ( map_args ) states = tree_map ( lambda x : jnp . swapaxes ( x , 0 , 1 ) , states ) states_flat = tree_map ( lambda x : jnp . reshape ( x , ( - 1 , ) + x . shape [ 2 : ] ) , states ) self . _last_state = last_state self . _states = states self . _states_flat = states_flat self . _set_collection_params ( )
short_docstring: Run the MCMC samplers and collect samples.
short_docstring_tokens: Run the MCMC samplers and collect samples .
comment: """
          Run the MCMC samplers and collect samples.
  
          :param random.PRNGKey rng_key: Random number generator key to be used for the sampling.
              For multi-chains, a batch of `num_chains` keys can be supplied. If `rng_key`
              does not have batch_size, it will be split in to a batch of `num_chains` keys.
          :param args: Arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init` method.
              These are typically the arguments needed by the `model`.
          :param extra_fields: Extra fields (aside from `"z"`, `"diverging"`) to be collected
              during the MCMC run. Note that subfields can be accessed using dots, e.g.
              `"adapt_state.step_size"` can be used to collect step sizes at each step.
          :type extra_fields: tuple or list of str
          :param init_params: Initial parameters to begin sampling. The type must be consistent
              with the input type to `potential_fn`.
          :param kwargs: Keyword arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init`
              method. These are typically the keyword arguments needed by the `model`.
  
          .. note:: jax allows python code to continue even when the compiled code has not finished yet.
              This can cause troubles when trying to profile the code for speed.
              See https://jax.readthedocs.io/en/latest/async_dispatch.html and
              https://jax.readthedocs.io/en/latest/profiling.html for pointers on profiling jax programs.
          """$SEP$# swap num_samples x num_chains to num_chains x num_samples