class_docstr: The variance exploding stochastic differential equation (SDE) scheduler.
  
  For more information, see the original paper: https://arxiv.org/abs/2011.13456
  
  [`~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.
      snr (`float`):
          coefficient weighting the step from the model_output sample (from the network) to the random noise.
      sigma_min (`float`):
              initial noise scale for sigma sequence in sampling procedure. The minimum sigma should mirror the
              distribution of the data.
      sigma_max (`float`): maximum value used for the range of continuous timesteps passed into the model.
      sampling_eps (`float`): the end value of sampling, where timesteps decrease progressively from 1 to
      epsilon.
      correct_steps (`int`): number of correction steps performed on a produced sample.
class_docstr_tok: ["The","variance","exploding","stochastic","differential","equation","(","SDE",")","scheduler",".","For","more","information",",","see","the","original","paper",":","https",":","\/\/arxiv.org\/abs\/2011.13456","[","`","~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",".","snr","(","`","float","`",")",":","coefficient","weighting","the","step","from","the","model_output","sample","(","from","the","network",")","to","the","random","noise",".","sigma_min","(","`","float","`",")",":","initial","noise","scale","for","sigma","sequence","in","sampling","procedure",".","The","minimum","sigma","should","mirror","the","distribution","of","the","data",".","sigma_max","(","`","float","`",")",":","maximum","value","used","for","the","range","of","continuous","timesteps","passed","into","the","model",".","sampling_eps","(","`","float","`",")",":","the","end","value","of","sampling",",","where","timesteps","decrease","progressively","from","1","to","epsilon",".","correct_steps","(","`","int","`",")",":","number","of","correction","steps","performed","on","a","produced","sample","."]
docstr: Correct the predicted sample based on the output model_output of the network. This is often run repeatedly
  after making the prediction for the previous timestep.
  
  Args:
      state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance.
      model_output (`jnp.ndarray`): direct output from learned diffusion model.
      sample (`jnp.ndarray`):
          current instance of sample being created by diffusion process.
      generator: random number generator.
      return_dict (`bool`): option for returning tuple rather than FlaxSdeVeOutput class
  
  Returns:
      [`FlaxSdeVeOutput`] or `tuple`: [`FlaxSdeVeOutput`] if `return_dict` is True, otherwise a `tuple`. When
      returning a tuple, the first element is the sample tensor.
docstr_tok: ["Correct","the","predicted","sample","based","on","the","output","model_output","of","the","network",".","This","is","often","run","repeatedly","after","making","the","prediction","for","the","previous","timestep",".","Args",":","state","(","`","ScoreSdeVeSchedulerState","`",")",":","the","`","FlaxScoreSdeVeScheduler","`","state","data","class","instance",".","model_output","(","`","jnp.ndarray","`",")",":","direct","output","from","learned","diffusion","model",".","sample","(","`","jnp.ndarray","`",")",":","current","instance","of","sample","being","created","by","diffusion","process",".","generator",":","random","number","generator",".","return_dict","(","`","bool","`",")",":","option","for","returning","tuple","rather","than","FlaxSdeVeOutput","class","Returns",":","[","`","FlaxSdeVeOutput","`","]","or","`","tuple","`",":","[","`","FlaxSdeVeOutput","`","]","if","`","return_dict","`","is","True",",","otherwise","a","`","tuple","`",".","When","returning","a","tuple",",","the","first","element","is","the","sample","tensor","."]
code: def step_correct(
      self,
      state: ScoreSdeVeSchedulerState,
      model_output: jnp.ndarray,
      sample: jnp.ndarray,
      key: jax.Array,
      return_dict: bool = True,
  ) -> Union[FlaxSdeVeOutput, Tuple]:
      """
      Correct the predicted sample based on the output model_output of the network. This is often run repeatedly
      after making the prediction for the previous timestep.
  
      Args:
          state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance.
          model_output (`jnp.ndarray`): direct output from learned diffusion model.
          sample (`jnp.ndarray`):
              current instance of sample being created by diffusion process.
          generator: random number generator.
          return_dict (`bool`): option for returning tuple rather than FlaxSdeVeOutput class
  
      Returns:
          [`FlaxSdeVeOutput`] or `tuple`: [`FlaxSdeVeOutput`] if `return_dict` is True, otherwise a `tuple`. When
          returning a tuple, the first element is the sample tensor.
  
      """
      if state.timesteps is None:
          raise ValueError(
              "`state.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
          )
  
      # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
      # sample noise for correction
      key = random.split(key, num=1)
      noise = random.normal(key=key, shape=sample.shape)
  
      # compute step size from the model_output, the noise, and the snr
      grad_norm = jnp.linalg.norm(model_output)
      noise_norm = jnp.linalg.norm(noise)
      step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
      step_size = step_size * jnp.ones(sample.shape[0])
  
      # compute corrected sample: model_output term and noise term
      step_size = step_size.flatten()
      step_size = broadcast_to_shape_from_left(step_size, sample.shape)
      prev_sample_mean = sample + step_size * model_output
      prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise
  
      if not return_dict:
          return (prev_sample, state)
  
      return FlaxSdeVeOutput(prev_sample=prev_sample, state=state)
  
imports: from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import flax
  import jax
  import jax.numpy
  from jax import random
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left
label_desc: Use image node_id 8 for calling the FlaxScoreSdeVeScheduler obj's underlying member method code with example usage: obj.step_correct(state, model_output, sample, key, return_dict) and returns: FlaxSdeVeOutput, prev_sample,  state