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: Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple
  times to approximate the solution.
  
  Args:
      state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance.
      model_output (`jnp.ndarray`): direct output from learned diffusion model.
      timestep (`int`): current discrete timestep in the diffusion chain.
      sample (`jnp.ndarray`):
          current instance of sample being created by diffusion process.
      return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class
  
  Returns:
      [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a
      `tuple`. When returning a tuple, the first element is the sample tensor.
docstr_tok: ["Step","function","propagating","the","sample","with","the","linear","multi-step","method",".","This","has","one","forward","pass","with","multiple","times","to","approximate","the","solution",".","Args",":","state","(","`","PNDMSchedulerState","`",")",":","the","`","FlaxPNDMScheduler","`","state","data","class","instance",".","model_output","(","`","jnp.ndarray","`",")",":","direct","output","from","learned","diffusion","model",".","timestep","(","`","int","`",")",":","current","discrete","timestep","in","the","diffusion","chain",".","sample","(","`","jnp.ndarray","`",")",":","current","instance","of","sample","being","created","by","diffusion","process",".","return_dict","(","`","bool","`",")",":","option","for","returning","tuple","rather","than","FlaxPNDMSchedulerOutput","class","Returns",":","[","`","FlaxPNDMSchedulerOutput","`","]","or","`","tuple","`",":","[","`","FlaxPNDMSchedulerOutput","`","]","if","`","return_dict","`","is","True",",","otherwise","a","`","tuple","`",".","When","returning","a","tuple",",","the","first","element","is","the","sample","tensor","."]
code: def step_plms(
      self,
      state: PNDMSchedulerState,
      model_output: jnp.ndarray,
      timestep: int,
      sample: jnp.ndarray,
  ) -> Union[FlaxPNDMSchedulerOutput, Tuple]:
      """
      Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple
      times to approximate the solution.
  
      Args:
          state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance.
          model_output (`jnp.ndarray`): direct output from learned diffusion model.
          timestep (`int`): current discrete timestep in the diffusion chain.
          sample (`jnp.ndarray`):
              current instance of sample being created by diffusion process.
          return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class
  
      Returns:
          [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a
          `tuple`. When returning a tuple, the first element is the sample tensor.
  
      """
  
      if state.num_inference_steps is None:
          raise ValueError(
              "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
          )
  
      # NOTE: There is no way to check in the jitted runtime if the prk mode was ran before
  
      prev_timestep = (
          timestep
          - self.config.num_train_timesteps // state.num_inference_steps
      )
      prev_timestep = jnp.where(prev_timestep > 0, prev_timestep, 0)
  
      # Reference:
      # if state.counter != 1:
      #     state.ets.append(model_output)
      # else:
      #     prev_timestep = timestep
      #     timestep = timestep + self.config.num_train_timesteps // state.num_inference_steps
  
      prev_timestep = jnp.where(
          state.counter == 1, timestep, prev_timestep
      )
      timestep = jnp.where(
          state.counter == 1,
          timestep
          + self.config.num_train_timesteps
          // state.num_inference_steps,
          timestep,
      )
  
      # Reference:
      # if len(state.ets) == 1 and state.counter == 0:
      #     model_output = model_output
      #     state.cur_sample = sample
      # elif len(state.ets) == 1 and state.counter == 1:
      #     model_output = (model_output + state.ets[-1]) / 2
      #     sample = state.cur_sample
      #     state.cur_sample = None
      # elif len(state.ets) == 2:
      #     model_output = (3 * state.ets[-1] - state.ets[-2]) / 2
      # elif len(state.ets) == 3:
      #     model_output = (23 * state.ets[-1] - 16 * state.ets[-2] + 5 * state.ets[-3]) / 12
      # else:
      #     model_output = (1 / 24) * (55 * state.ets[-1] - 59 * state.ets[-2] + 37 * state.ets[-3] - 9 * state.ets[-4])
  
      state = state.replace(
          ets=jax.lax.select(
              state.counter != 1,
              state.ets.at[0:3]
              .set(state.ets[1:4])
              .at[3]
              .set(model_output),  # counter != 1
              state.ets,  # counter 1
          ),
          cur_sample=jax.lax.select(
              state.counter != 1,
              sample,  # counter != 1
              state.cur_sample,  # counter 1
          ),
      )
  
      state = state.replace(
          cur_model_output=jax.lax.select_n(
              jnp.clip(state.counter, 0, 4),
              model_output,  # counter 0
              (model_output + state.ets[-1]) / 2,  # counter 1
              (3 * state.ets[-1] - state.ets[-2]) / 2,  # counter 2
              (
                  23 * state.ets[-1]
                  - 16 * state.ets[-2]
                  + 5 * state.ets[-3]
              )
              / 12,  # counter 3
              (1 / 24)
              * (
                  55 * state.ets[-1]
                  - 59 * state.ets[-2]
                  + 37 * state.ets[-3]
                  - 9 * state.ets[-4]
              ),  # counter >= 4
          ),
      )
  
      sample = state.cur_sample
      model_output = state.cur_model_output
      prev_sample = self._get_prev_sample(
          state, sample, timestep, prev_timestep, model_output
      )
      state = state.replace(counter=state.counter + 1)
  
      return (prev_sample, state)
  
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 8 for calling the FlaxPNDMScheduler obj's underlying member method code with example usage: obj.step_plms(state, model_output, timestep, sample) and returns: prev_sample,  state