class_docstr: Stochastic sampling from Karras et al. [1] tailored to the Variance-Expanding (VE) models [2]. Use Algorithm 2 and
  the VE column of Table 1 from [1] for reference.
  
  [1] Karras, Tero, et al. "Elucidating the Design Space of Diffusion-Based Generative Models."
  https://arxiv.org/abs/2206.00364 [2] Song, Yang, et al. "Score-based generative modeling through stochastic
  differential equations." 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.
  
  For more details on the parameters, see the original paper's Appendix E.: "Elucidating the Design Space of
  Diffusion-Based Generative Models." https://arxiv.org/abs/2206.00364. The grid search values used to find the
  optimal {s_noise, s_churn, s_min, s_max} for a specific model are described in Table 5 of the paper.
  
  Args:
      sigma_min (`float`): minimum noise magnitude
      sigma_max (`float`): maximum noise magnitude
      s_noise (`float`): the amount of additional noise to counteract loss of detail during sampling.
          A reasonable range is [1.000, 1.011].
      s_churn (`float`): the parameter controlling the overall amount of stochasticity.
          A reasonable range is [0, 100].
      s_min (`float`): the start value of the sigma range where we add noise (enable stochasticity).
          A reasonable range is [0, 10].
      s_max (`float`): the end value of the sigma range where we add noise.
          A reasonable range is [0.2, 80].
code: def __init__(
      self,
      sigma_min: float = 0.02,
      sigma_max: float = 100,
      s_noise: float = 1.007,
      s_churn: float = 80,
      s_min: float = 0.05,
      s_max: float = 50,
  ):
      pass
  
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 ..utils import BaseOutput
  from .scheduling_utils_flax import FlaxSchedulerMixin
label_desc: Use image node_id 2 to create a new FlaxKarrasVeScheduler object from inherited base classes: FlaxSchedulerMixin, ConfigMixin with example: obj = FlaxKarrasVeScheduler(sigma_min, sigma_max, s_noise, s_churn, s_min, s_max)