class_docstr: DPM-Solver (and the improved version DPM-Solver++) is a fast dedicated high-order solver for diffusion ODEs with
  the convergence order guarantee. Empirically, sampling by DPM-Solver with only 20 steps can generate high-quality
  samples, and it can generate quite good samples even in only 10 steps.
  
  For more details, see the original paper: https://arxiv.org/abs/2206.00927 and https://arxiv.org/abs/2211.01095
  
  Currently, we support the multistep DPM-Solver for both noise prediction models and data prediction models. We
  recommend to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling.
  
  We also support the "dynamic thresholding" method in Imagen (https://arxiv.org/abs/2205.11487). For pixel-space
  diffusion models, you can set both `algorithm_type="dpmsolver++"` and `thresholding=True` to use the dynamic
  thresholding. Note that the thresholding method is unsuitable for latent-space diffusion models (such as
  stable-diffusion).
  
  [`~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/2206.00927 and https://arxiv.org/abs/2211.01095
  
  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 (`np.ndarray`, optional):
          option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.
      solver_order (`int`, default `2`):
          the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided
          sampling, and `solver_order=3` for unconditional sampling.
      prediction_type (`str`, default `epsilon`):
          indicates whether the model predicts the noise (epsilon), or the data / `x0`. One of `epsilon`, `sample`,
          or `v-prediction`.
      thresholding (`bool`, default `False`):
          whether to use the "dynamic thresholding" method (introduced by Imagen, https://arxiv.org/abs/2205.11487).
          For pixel-space diffusion models, you can set both `algorithm_type=dpmsolver++` and `thresholding=True` to
          use the dynamic thresholding. Note that the thresholding method is unsuitable for latent-space diffusion
          models (such as stable-diffusion).
      dynamic_thresholding_ratio (`float`, default `0.995`):
          the ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen
          (https://arxiv.org/abs/2205.11487).
      sample_max_value (`float`, default `1.0`):
          the threshold value for dynamic thresholding. Valid only when `thresholding=True` and
          `algorithm_type="dpmsolver++`.
      algorithm_type (`str`, default `dpmsolver++`):
          the algorithm type for the solver. Either `dpmsolver` or `dpmsolver++`. The `dpmsolver` type implements the
          algorithms in https://arxiv.org/abs/2206.00927, and the `dpmsolver++` type implements the algorithms in
          https://arxiv.org/abs/2211.01095. We recommend to use `dpmsolver++` with `solver_order=2` for guided
          sampling (e.g. stable-diffusion).
      solver_type (`str`, default `midpoint`):
          the solver type for the second-order solver. Either `midpoint` or `heun`. The solver type slightly affects
          the sample quality, especially for small number of steps. We empirically find that `midpoint` solvers are
          slightly better, so we recommend to use the `midpoint` type.
      lower_order_final (`bool`, default `True`):
          whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically
          find this trick can stabilize the sampling of DPM-Solver for steps < 15, especially for steps <= 10.
      timestep_spacing (`str`, defaults to `"linspace"`):
          The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
          Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
      dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
          the `dtype` used for params and computation.
class_docstr_tok: ["DPM-Solver","(","and","the","improved","version","DPM-Solver++",")","is","a","fast","dedicated","high-order","solver","for","diffusion","ODEs","with","the","convergence","order","guarantee",".","Empirically",",","sampling","by","DPM-Solver","with","only","20","steps","can","generate","high-quality","samples",",","and","it","can","generate","quite","good","samples","even","in","only","10","steps",".","For","more","details",",","see","the","original","paper",":","https",":","\/\/arxiv.org\/abs\/2206.00927","and","https",":","\/\/arxiv.org\/abs\/2211.01095","Currently",",","we","support","the","multistep","DPM-Solver","for","both","noise","prediction","models","and","data","prediction","models",".","We","recommend","to","use","`","solver_order=2","`","for","guided","sampling",",","and","`","solver_order=3","`","for","unconditional","sampling",".","We","also","support","the","``","dynamic","thresholding","''","method","in","Imagen","(","https",":","\/\/arxiv.org\/abs\/2205.11487",")",".","For","pixel-space","diffusion","models",",","you","can","set","both","`","algorithm_type=","''","dpmsolver++","''","`","and","`","thresholding=True","`","to","use","the","dynamic","thresholding",".","Note","that","the","thresholding","method","is","unsuitable","for","latent-space","diffusion","models","(","such","as","stable-diffusion",")",".","[","`","~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\/2206.00927","and","https",":","\/\/arxiv.org\/abs\/2211.01095","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","(","`","np.ndarray","`",",","optional",")",":","option","to","pass","an","array","of","betas","directly","to","the","constructor","to","bypass","`","beta_start","`",",","`","beta_end","`","etc",".","solver_order","(","`","int","`",",","default","`","2","`",")",":","the","order","of","DPM-Solver",";","can","be","`","1","`","or","`","2","`","or","`","3","`",".","We","recommend","to","use","`","solver_order=2","`","for","guided","sampling",",","and","`","solver_order=3","`","for","unconditional","sampling",".","prediction_type","(","`","str","`",",","default","`","epsilon","`",")",":","indicates","whether","the","model","predicts","the","noise","(","epsilon",")",",","or","the","data","\/","`","x0","`",".","One","of","`","epsilon","`",",","`","sample","`",",","or","`","v-prediction","`",".","thresholding","(","`","bool","`",",","default","`","False","`",")",":","whether","to","use","the","``","dynamic","thresholding","''","method","(","introduced","by","Imagen",",","https",":","\/\/arxiv.org\/abs\/2205.11487",")",".","For","pixel-space","diffusion","models",",","you","can","set","both","`","algorithm_type=dpmsolver++","`","and","`","thresholding=True","`","to","use","the","dynamic","thresholding",".","Note","that","the","thresholding","method","is","unsuitable","for","latent-space","diffusion","models","(","such","as","stable-diffusion",")",".","dynamic_thresholding_ratio","(","`","float","`",",","default","`","0.995","`",")",":","the","ratio","for","the","dynamic","thresholding","method",".","Default","is","`","0.995","`",",","the","same","as","Imagen","(","https",":","\/\/arxiv.org\/abs\/2205.11487",")",".","sample_max_value","(","`","float","`",",","default","`","1.0","`",")",":","the","threshold","value","for","dynamic","thresholding",".","Valid","only","when","`","thresholding=True","`","and","`","algorithm_type=","''","dpmsolver++","`",".","algorithm_type","(","`","str","`",",","default","`","dpmsolver++","`",")",":","the","algorithm","type","for","the","solver",".","Either","`","dpmsolver","`","or","`","dpmsolver++","`",".","The","`","dpmsolver","`","type","implements","the","algorithms","in","https",":","\/\/arxiv.org\/abs\/2206.00927",",","and","the","`","dpmsolver++","`","type","implements","the","algorithms","in","https",":","\/\/arxiv.org\/abs\/2211.01095",".","We","recommend","to","use","`","dpmsolver++","`","with","`","solver_order=2","`","for","guided","sampling","(","e.g",".","stable-diffusion",")",".","solver_type","(","`","str","`",",","default","`","midpoint","`",")",":","the","solver","type","for","the","second-order","solver",".","Either","`","midpoint","`","or","`","heun","`",".","The","solver","type","slightly","affects","the","sample","quality",",","especially","for","small","number","of","steps",".","We","empirically","find","that","`","midpoint","`","solvers","are","slightly","better",",","so","we","recommend","to","use","the","`","midpoint","`","type",".","lower_order_final","(","`","bool","`",",","default","`","True","`",")",":","whether","to","use","lower-order","solvers","in","the","final","steps",".","Only","valid","for","<","15","inference","steps",".","We","empirically","find","this","trick","can","stabilize","the","sampling","of","DPM-Solver","for","steps","<","15",",","especially","for","steps","<","=","10.","timestep_spacing","(","`","str","`",",","defaults","to","`","``","linspace","''","`",")",":","The","way","the","timesteps","should","be","scaled",".","Refer","to","Table","2","of","the","[","Common","Diffusion","Noise","Schedules","and","Sample","Steps","are","Flawed","]","(","https",":","\/\/huggingface.co\/papers\/2305.08891",")","for","more","information",".","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 (`DPMSolverMultistepSchedulerState`):
          the `FlaxDPMSolverMultistepScheduler` 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","(","`","DPMSolverMultistepSchedulerState","`",")",":","the","`","FlaxDPMSolverMultistepScheduler","`","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: DPMSolverMultistepSchedulerState,
      num_inference_steps: int,
      shape: Tuple,
  ) -> DPMSolverMultistepSchedulerState:
      """
      Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
  
      Args:
          state (`DPMSolverMultistepSchedulerState`):
              the `FlaxDPMSolverMultistepScheduler` 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.
      """
      last_timestep = self.config.num_train_timesteps
      if self.config.timestep_spacing == "linspace":
          timesteps = (
              jnp.linspace(
                  0, last_timestep - 1, num_inference_steps + 1
              )
              .round()[::-1][:-1]
              .astype(jnp.int32)
          )
      elif self.config.timestep_spacing == "leading":
          step_ratio = last_timestep // (num_inference_steps + 1)
          # creates integer timesteps by multiplying by ratio
          # casting to int to avoid issues when num_inference_step is power of 3
          timesteps = (
              (jnp.arange(0, num_inference_steps + 1) * step_ratio)
              .round()[::-1][:-1]
              .copy()
              .astype(jnp.int32)
          )
          timesteps += self.config.steps_offset
      elif self.config.timestep_spacing == "trailing":
          step_ratio = (
              self.config.num_train_timesteps / num_inference_steps
          )
          # creates integer timesteps by multiplying by ratio
          # casting to int to avoid issues when num_inference_step is power of 3
          timesteps = (
              jnp.arange(last_timestep, 0, -step_ratio)
              .round()
              .copy()
              .astype(jnp.int32)
          )
          timesteps -= 1
      else:
          raise ValueError(
              f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
          )
  
      # initial running values
  
      model_outputs = jnp.zeros(
          (self.config.solver_order,) + shape, dtype=self.dtype
      )
      lower_order_nums = jnp.int32(0)
      prev_timestep = jnp.int32(-1)
      cur_sample = jnp.zeros(shape, dtype=self.dtype)
  
      return state.replace(
          num_inference_steps=num_inference_steps,
          timesteps=timesteps,
          model_outputs=model_outputs,
          lower_order_nums=lower_order_nums,
          prev_timestep=prev_timestep,
          cur_sample=cur_sample,
      )
  
imports: from dataclasses import dataclass
  from typing import List, 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 FlaxDPMSolverMultistepScheduler obj's underlying member method code with example usage: obj.set_timesteps(state, num_inference_steps, shape) and returns: state