class_docstr: Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by
  Katherine Crowson:
  https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181
  
  [`~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.
  
  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` or `scaled_linear`.
      trained_betas (`jnp.ndarray`, optional):
          option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.
      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 timesteps used for the diffusion chain. Supporting function to be run before inference.
  
  Args:
      state (`LMSDiscreteSchedulerState`):
          the `FlaxLMSDiscreteScheduler` state data class instance.
      num_inference_steps (`int`):
          the number of diffusion steps used when generating samples with a pre-trained model.
docstr_tok: ["Sets","the","timesteps","used","for","the","diffusion","chain",".","Supporting","function","to","be","run","before","inference",".","Args",":","state","(","`","LMSDiscreteSchedulerState","`",")",":","the","`","FlaxLMSDiscreteScheduler","`","state","data","class","instance",".","num_inference_steps","(","`","int","`",")",":","the","number","of","diffusion","steps","used","when","generating","samples","with","a","pre-trained","model","."]
code: def set_timesteps(
      self,
      state: LMSDiscreteSchedulerState,
      num_inference_steps: int,
      shape: Tuple = (),
  ) -> LMSDiscreteSchedulerState:
      """
      Sets the timesteps used for the diffusion chain. Supporting function to be run before inference.
  
      Args:
          state (`LMSDiscreteSchedulerState`):
              the `FlaxLMSDiscreteScheduler` state data class instance.
          num_inference_steps (`int`):
              the number of diffusion steps used when generating samples with a pre-trained model.
      """
  
      timesteps = jnp.linspace(
          self.config.num_train_timesteps - 1,
          0,
          num_inference_steps,
          dtype=self.dtype,
      )
  
      low_idx = jnp.floor(timesteps).astype(jnp.int32)
      high_idx = jnp.ceil(timesteps).astype(jnp.int32)
  
      frac = jnp.mod(timesteps, 1.0)
  
      sigmas = (
          (1 - state.common.alphas_cumprod)
          / state.common.alphas_cumprod
      ) ** 0.5
      sigmas = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
      sigmas = jnp.concatenate(
          [sigmas, jnp.array([0.0], dtype=self.dtype)]
      )
  
      timesteps = timesteps.astype(jnp.int32)
  
      # initial running values
      derivatives = jnp.zeros((0,) + shape, dtype=self.dtype)
  
      return state.replace(
          timesteps=timesteps,
          sigmas=sigmas,
          num_inference_steps=num_inference_steps,
          derivatives=derivatives,
      )
  
imports: from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import flax
  import jax.numpy
  from scipy import integrate
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import CommonSchedulerState, FlaxKarrasDiffusionSchedulers, FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left
label_desc: Use image node_id 6 for calling the FlaxLMSDiscreteScheduler obj's underlying member method code with example usage: obj.set_timesteps(state, num_inference_steps, shape) and returns: state