class_docstr: Pseudo numerical methods for diffusion models (PNDM) proposes using more advanced ODE integration techniques,
  namely Runge-Kutta method and a linear multi-step method.
  
  [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
  function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
  [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and
  [`~SchedulerMixin.from_pretrained`] functions.
  
  For more details, see the original paper: https://arxiv.org/abs/2202.09778
  
  Args:
      num_train_timesteps (`int`): number of diffusion steps used to train the model.
      beta_start (`float`): the starting `beta` value of inference.
      beta_end (`float`): the final `beta` value.
      beta_schedule (`str`):
          the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
          `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
      trained_betas (`jnp.ndarray`, optional):
          option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.
      skip_prk_steps (`bool`):
          allows the scheduler to skip the Runge-Kutta steps that are defined in the original paper as being required
          before plms steps; defaults to `False`.
      set_alpha_to_one (`bool`, default `False`):
          each diffusion step uses the value of alphas product at that step and at the previous one. For the final
          step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
          otherwise it uses the value of alpha at step 0.
      steps_offset (`int`, default `0`):
          an offset added to the inference steps. You can use a combination of `offset=1` and
          `set_alpha_to_one=False`, to make the last step use step 0 for the previous alpha product, as done in
          stable diffusion.
      prediction_type (`str`, default `epsilon`, optional):
          prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion
          process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4
          https://imagen.research.google/video/paper.pdf)
      dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
          the `dtype` used for params and computation.
docstr: Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
  
  Args:
      state (`PNDMSchedulerState`):
          the `FlaxPNDMScheduler` state data class instance.
      num_inference_steps (`int`):
          the number of diffusion steps used when generating samples with a pre-trained model.
      shape (`Tuple`):
          the shape of the samples to be generated.
docstr_tok: ["Sets","the","discrete","timesteps","used","for","the","diffusion","chain",".","Supporting","function","to","be","run","before","inference",".","Args",":","state","(","`","PNDMSchedulerState","`",")",":","the","`","FlaxPNDMScheduler","`","state","data","class","instance",".","num_inference_steps","(","`","int","`",")",":","the","number","of","diffusion","steps","used","when","generating","samples","with","a","pre-trained","model",".","shape","(","`","Tuple","`",")",":","the","shape","of","the","samples","to","be","generated","."]
code: def set_timesteps(
      self,
      state: PNDMSchedulerState,
      num_inference_steps: int,
      shape: Tuple,
  ) -> PNDMSchedulerState:
      """
      Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
  
      Args:
          state (`PNDMSchedulerState`):
              the `FlaxPNDMScheduler` state data class instance.
          num_inference_steps (`int`):
              the number of diffusion steps used when generating samples with a pre-trained model.
          shape (`Tuple`):
              the shape of the samples to be generated.
      """
  
      step_ratio = (
          self.config.num_train_timesteps // num_inference_steps
      )
      # creates integer timesteps by multiplying by ratio
      # rounding to avoid issues when num_inference_step is power of 3
      _timesteps = (
          jnp.arange(0, num_inference_steps) * step_ratio
      ).round() + self.config.steps_offset
  
      if self.config.skip_prk_steps:
          # for some models like stable diffusion the prk steps can/should be skipped to
          # produce better results. When using PNDM with `self.config.skip_prk_steps` the implementation
          # is based on crowsonkb's PLMS sampler implementation: https://github.com/CompVis/latent-diffusion/pull/51
  
          prk_timesteps = jnp.array([], dtype=jnp.int32)
          plms_timesteps = jnp.concatenate(
              [_timesteps[:-1], _timesteps[-2:-1], _timesteps[-1:]]
          )[::-1]
  
      else:
          prk_timesteps = _timesteps[-self.pndm_order :].repeat(
              2
          ) + jnp.tile(
              jnp.array(
                  [
                      0,
                      self.config.num_train_timesteps
                      // num_inference_steps
                      // 2,
                  ],
                  dtype=jnp.int32,
              ),
              self.pndm_order,
          )
  
          prk_timesteps = (prk_timesteps[:-1].repeat(2)[1:-1])[::-1]
          plms_timesteps = _timesteps[:-3][::-1]
  
      timesteps = jnp.concatenate([prk_timesteps, plms_timesteps])
  
      # initial running values
  
      cur_model_output = jnp.zeros(shape, dtype=self.dtype)
      counter = jnp.int32(0)
      cur_sample = jnp.zeros(shape, dtype=self.dtype)
      ets = jnp.zeros((4,) + shape, dtype=self.dtype)
  
      return state.replace(
          timesteps=timesteps,
          num_inference_steps=num_inference_steps,
          prk_timesteps=prk_timesteps,
          plms_timesteps=plms_timesteps,
          cur_model_output=cur_model_output,
          counter=counter,
          cur_sample=cur_sample,
          ets=ets,
      )
  
imports: from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import flax
  import jax
  import jax.numpy
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import CommonSchedulerState, FlaxKarrasDiffusionSchedulers, FlaxSchedulerMixin, FlaxSchedulerOutput, add_noise_common
label_desc: Use image node_id 4 for calling the FlaxPNDMScheduler obj's underlying member method code with example usage: obj.set_timesteps(state, num_inference_steps, shape) and returns: state