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].
class_docstr_tok: ["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","]","."]
docstr: Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a
  higher noise level sigma_hat = sigma_i + gamma_i*sigma_i.
  
  TODO Args:
docstr_tok: ["Explicit","Langevin-like","``","churn","''","step","of","adding","noise","to","the","sample","according","to","a","factor","gamma_i","\u2265","0","to","reach","a","higher","noise","level","sigma_hat","=","sigma_i","+","gamma_i","*","sigma_i",".","TODO","Args",":"]
code: def add_noise_to_input(
      self,
      state: KarrasVeSchedulerState,
      sample: jnp.ndarray,
      sigma: float,
      key: jax.Array,
  ) -> Tuple[jnp.ndarray, float]:
      """
      Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a
      higher noise level sigma_hat = sigma_i + gamma_i*sigma_i.
  
      TODO Args:
      """
      if self.config.s_min <= sigma <= self.config.s_max:
          gamma = min(
              self.config.s_churn / state.num_inference_steps,
              2**0.5 - 1,
          )
      else:
          gamma = 0
  
      # sample eps ~ N(0, S_noise^2 * I)
      key = random.split(key, num=1)
      eps = self.config.s_noise * random.normal(
          key=key, shape=sample.shape
      )
      sigma_hat = sigma + gamma * sigma
      sample_hat = sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps)
  
      return sample_hat, sigma_hat
  
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 5 for calling the FlaxKarrasVeScheduler obj's underlying member method code with example usage: obj.add_noise_to_input(state, sample, sigma, key) and returns: sample_hat,  sigma_hat