code: 
  
  
  # DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
  
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  
  import flax
  import jax.numpy as jnp
  from jax import random
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left
  
  @flax.struct.dataclass
  class 						UpperCAmelCase__   :
  
  
  
       '''simple docstring'''
       UpperCAmelCase_         =  None
       UpperCAmelCase_         =  None
       UpperCAmelCase_         =  None
  
  
  
  
       @classmethod
       def   lowerCAmelCase_   ( cls       :       List[Any]		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              return cls()
  
  
  
  
  @dataclass
  class 						UpperCAmelCase__   (      UpperCAmelCase__		):
  
  
  
       '''simple docstring'''
       UpperCAmelCase_         =  42
       UpperCAmelCase_         =  42
       UpperCAmelCase_         =  None
  
  
  
  
  class 						UpperCAmelCase__   (      UpperCAmelCase__      ,						UpperCAmelCase__		):
  
  
  
       '''simple docstring'''
  
  
  
  
       @property
       def   lowerCAmelCase_   ( self       :       Dict		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              return True
  
  
  
  
       @register_to_config
       def __init__( self       :       str	,      UpperCamelCase       :       Optional[Any] = 20_00	,      UpperCamelCase       :       Optional[Any] = 0.15	,      UpperCamelCase       :       Optional[int] = 0.01	,      UpperCamelCase       :       Any = 13_48.0	,      UpperCamelCase       :       List[Any] = 1E-5	,      UpperCamelCase       :       Dict = 1	,      ):
  
  
  
  
  
              """simple docstring"""
  
  
  
              pass
  
  
  
  
       def   lowerCAmelCase_   ( self       :       Union[str, Any]		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              _lowercase					:				Tuple           =			ScoreSdeVeSchedulerState.create()
              return self.set_sigmas(
                  lowerCamelCase__	,      self.config.num_train_timesteps	,      self.config.sigma_min	,      self.config.sigma_max	,      self.config.sampling_eps	,      )
  
  
  
  
       def   lowerCAmelCase_   ( self       :       int	,      UpperCamelCase       :       Optional[Any]	,      UpperCamelCase       :       Dict	,      UpperCamelCase       :       str = ()	,      UpperCamelCase       :       Tuple = None		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              _lowercase					:				str           =			sampling_eps if sampling_eps is not None else self.config.sampling_eps
  
              _lowercase					:				List[str]           =			jnp.linspace(1	,      lowerCamelCase__	,      lowerCamelCase__		)
              return state.replace(timesteps=lowerCamelCase__		)
  
  
  
  
       def   lowerCAmelCase_   ( self       :       List[Any]	,      UpperCamelCase       :       Optional[int]	,      UpperCamelCase       :       Optional[Any]	,      UpperCamelCase       :       Optional[int] = None	,      UpperCamelCase       :       str = None	,      UpperCamelCase       :       Optional[int] = None	,      ):
  
  
  
  
  
              """simple docstring"""
  
  
  
              _lowercase					:				Dict           =			sigma_min if sigma_min is not None else self.config.sigma_min
              _lowercase					:				Dict           =			sigma_max if sigma_max is not None else self.config.sigma_max
              _lowercase					:				int           =			sampling_eps if sampling_eps is not None else self.config.sampling_eps
              if state.timesteps is None:
                     _lowercase					:				Dict           =			self.set_timesteps(lowerCamelCase__	,      lowerCamelCase__	,      lowerCamelCase__		)
  
              _lowercase					:				Optional[int]           =			jnp.exp(jnp.linspace(jnp.log(lowerCamelCase__		)	,      jnp.log(lowerCamelCase__		)	,      lowerCamelCase__		)		)
              _lowercase					:				Union[str, Any]           =			jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps]		)
  
              return state.replace(discrete_sigmas=lowerCamelCase__	,      sigmas=lowerCamelCase__		)
  
  
  
  
       def   lowerCAmelCase_   ( self       :       List[str]	,      UpperCamelCase       :       Dict	,      UpperCamelCase       :       str	,      UpperCamelCase       :       Optional[Any]		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              return jnp.where(timesteps == 0	,      jnp.zeros_like(lowerCamelCase__		)	,      state.discrete_sigmas[timesteps - 1]		)
  
  
  
  
       def   lowerCAmelCase_   ( self       :       Dict	,      UpperCamelCase       :       Optional[Any]	,      UpperCamelCase       :       List[str]	,      UpperCamelCase       :       int	,      UpperCamelCase       :       List[str]	,      UpperCamelCase       :       Optional[int]	,      UpperCamelCase       :       Union[str, Any] = True	,      ):
  
  
  
  
  
              """simple docstring"""
  
  
  
              if state.timesteps is None:
                     raise ValueError(
                         '''`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'''		)
  
              _lowercase					:				Optional[int]           =			timestep * jnp.ones(
                  sample.shape[0]	,      )
              _lowercase					:				Dict           =			(timestep * (len(state.timesteps		) - 1)).long()
  
              _lowercase					:				Optional[int]           =			state.discrete_sigmas[timesteps]
              _lowercase					:				int           =			self.get_adjacent_sigma(lowerCamelCase__	,      lowerCamelCase__	,      lowerCamelCase__		)
              _lowercase					:				Optional[Any]           =			jnp.zeros_like(lowerCamelCase__		)
              _lowercase					:				int           =			(sigma**2 - adjacent_sigma**2) ** 0.5
  
              # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
              # also equation 47 shows the analog from SDE models to ancestral sampling methods
              _lowercase					:				Any           =			diffusion.flatten()
              _lowercase					:				Union[str, Any]           =			broadcast_to_shape_from_left(lowerCamelCase__	,      sample.shape		)
              _lowercase					:				Union[str, Any]           =			drift - diffusion**2 * model_output
  
              #  equation 6: sample noise for the diffusion term of
              _lowercase					:				Optional[Any]           =			random.split(lowerCamelCase__	,      num=1		)
              _lowercase					:				Union[str, Any]           =			random.normal(key=lowerCamelCase__	,      shape=sample.shape		)
              _lowercase					:				Dict           =			sample - drift  # subtract because `dt` is a small negative timestep
              # TODO is the variable diffusion the correct scaling term for the noise?
              _lowercase					:				Any           =			prev_sample_mean + diffusion * noise  # add impact of diffusion field g
  
              if not return_dict:
                     return (prev_sample, prev_sample_mean, state)
  
              return FlaxSdeVeOutput(prev_sample=lowerCamelCase__	,      prev_sample_mean=lowerCamelCase__	,      state=lowerCamelCase__		)
  
  
  
  
       def   lowerCAmelCase_   ( self       :       Optional[Any]	,      UpperCamelCase       :       List[str]	,      UpperCamelCase       :       List[Any]	,      UpperCamelCase       :       int	,      UpperCamelCase       :       List[str]	,      UpperCamelCase       :       Dict = True	,      ):
  
  
  
  
  
              """simple docstring"""
  
  
  
              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
              _lowercase					:				Tuple           =			random.split(lowerCamelCase__	,      num=1		)
              _lowercase					:				Any           =			random.normal(key=lowerCamelCase__	,      shape=sample.shape		)
  
              # compute step size from the model_output, the noise, and the snr
              _lowercase					:				Tuple           =			jnp.linalg.norm(lowerCamelCase__		)
              _lowercase					:				Any           =			jnp.linalg.norm(lowerCamelCase__		)
              _lowercase					:				Dict           =			(self.config.snr * noise_norm / grad_norm) ** 2 * 2
              _lowercase					:				Dict           =			step_size * jnp.ones(sample.shape[0]		)
  
              # compute corrected sample: model_output term and noise term
              _lowercase					:				List[Any]           =			step_size.flatten()
              _lowercase					:				str           =			broadcast_to_shape_from_left(lowerCamelCase__	,      sample.shape		)
              _lowercase					:				str           =			sample + step_size * model_output
              _lowercase					:				str           =			prev_sample_mean + ((step_size * 2) ** 0.5) * noise
  
              if not return_dict:
                     return (prev_sample, state)
  
              return FlaxSdeVeOutput(prev_sample=lowerCamelCase__	,      state=lowerCamelCase__		)
  
  
  
  
       def __len__( self       :       int		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              return self.config.num_train_timesteps
  
  
  
style_context: 
  
  
  import torch
  import torch.nn as nn
  from transformers.modeling_utils import ModuleUtilsMixin
  from transformers.models.ta.modeling_ta import (
      TaBlock,
      TaConfig,
      TaLayerNorm,
  )
  
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...models import ModelMixin
  
  class 						UpperCAmelCase__   (      A_      ,						A_      ,						A_		):
  
  
  
       '''simple docstring'''
  
  
  
  
       @register_to_config
       def __init__( self       :       Dict	,      UpperCamelCase       :       int	,      UpperCamelCase       :       int	,      UpperCamelCase       :       int	,      UpperCamelCase       :       float	,      UpperCamelCase       :       int	,      UpperCamelCase       :       int	,      UpperCamelCase       :       int	,      UpperCamelCase       :       int	,      UpperCamelCase       :       str	,      UpperCamelCase       :       bool = False	,      ):
  
  
  
  
  
              """simple docstring"""
  
  
  
              super().__init__()
  
              _lowercase					:				Any           =			nn.Linear(UpperCamelCase	,      UpperCamelCase	,      bias=UpperCamelCase		)
  
              _lowercase					:				List[Any]           =			nn.Embedding(UpperCamelCase	,      UpperCamelCase		)
              _lowercase					:				Any           =			False
  
              _lowercase					:				Tuple           =			nn.Dropout(p=UpperCamelCase		)
  
              _lowercase					:				Any           =			TaConfig(
                  d_model=UpperCamelCase	,      num_heads=UpperCamelCase	,      d_kv=UpperCamelCase	,      d_ff=UpperCamelCase	,      feed_forward_proj=UpperCamelCase	,      dropout_rate=UpperCamelCase	,      is_decoder=UpperCamelCase	,      is_encoder_decoder=UpperCamelCase	,      )
              _lowercase					:				Any           =			nn.ModuleList()
              for lyr_num in range(UpperCamelCase		):
                     _lowercase					:				Tuple           =			TaBlock(UpperCamelCase		)
                     self.encoders.append(UpperCamelCase		)
  
              _lowercase					:				List[Any]           =			TaLayerNorm(UpperCamelCase		)
              _lowercase					:				Optional[Any]           =			nn.Dropout(p=UpperCamelCase		)
  
  
  
  
       def   lowerCAmelCase_   ( self       :       int	,      UpperCamelCase       :       Optional[int]	,      UpperCamelCase       :       List[Any]		):
  
  
  
  
  
              """simple docstring"""
  
  
  
              _lowercase					:				Optional[int]           =			self.input_proj(UpperCamelCase		)
  
              # terminal relative positional encodings
              _lowercase					:				Optional[int]           =			encoder_inputs.shape[1]
              _lowercase					:				Optional[int]           =			torch.arange(UpperCamelCase	,      device=encoder_inputs.device		)
  
              _lowercase					:				Any           =			encoder_inputs_mask.sum(-1		)
              _lowercase					:				int           =			torch.roll(input_positions.unsqueeze(0		)	,      tuple(seq_lens.tolist()		)	,      dims=0		)
              x += self.position_encoding(UpperCamelCase		)
  
              _lowercase					:				Optional[int]           =			self.dropout_pre(UpperCamelCase		)
  
              # inverted the attention mask
              _lowercase					:				Optional[int]           =			encoder_inputs.size()
              _lowercase					:				List[Any]           =			self.get_extended_attention_mask(UpperCamelCase	,      UpperCamelCase		)
  
              for lyr in self.encoders:
                     _lowercase					:				Optional[int]           =			lyr(UpperCamelCase	,      UpperCamelCase		)[0]
              _lowercase					:				List[str]           =			self.layer_norm(UpperCamelCase		)
  
              return self.dropout_post(UpperCamelCase		), encoder_inputs_mask