


'''simple docstring'''







import torch

from diffusers import EulerAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class 		UpperCamelCase__    (       _UpperCamelCase		):



     """simple docstring"""






     __magic_name__				=			(EulerAncestralDiscreteScheduler,)
     __magic_name__				=			1_0







     def a		(       self					,		**snake_case__      ):

            '''simple docstring'''

            _lowerCAmelCase			:							List[str]      	=				{
                "num_train_timesteps": 1100,
                "beta_start": 0.0001,
                "beta_end": 0.02,
                "beta_schedule": "linear",
            }

            config.update(**__a      )
            return config







     def a		(       self      ):

            '''simple docstring'''

            for timesteps in [10, 50, 100, 1000]:
                   self.check_over_configs(num_train_timesteps=__a      )







     def a		(       self      ):

            '''simple docstring'''

            for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]					,		[0.0002, 0.002, 0.02]      ):
                   self.check_over_configs(beta_start=__a					,		beta_end=__a      )







     def a		(       self      ):

            '''simple docstring'''

            for schedule in ["linear", "scaled_linear"]:
                   self.check_over_configs(beta_schedule=__a      )







     def a		(       self      ):

            '''simple docstring'''

            for prediction_type in ["epsilon", "v_prediction"]:
                   self.check_over_configs(prediction_type=__a      )







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							int      	=				self.scheduler_classes[0]
            _lowerCAmelCase			:							Tuple      	=				self.get_scheduler_config()
            _lowerCAmelCase			:							Dict      	=				scheduler_class(**__a      )

            scheduler.set_timesteps(self.num_inference_steps      )

            _lowerCAmelCase			:							Any      	=				torch.manual_seed(0      )

            _lowerCAmelCase			:							Union[str, Any]      	=				self.dummy_model()
            _lowerCAmelCase			:							Dict      	=				self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
            _lowerCAmelCase			:							Dict      	=				sample.to(__a      )

            for i, t in enumerate(scheduler.timesteps      ):
                   _lowerCAmelCase			:							Optional[int]      	=				scheduler.scale_model_input(__a					,		__a      )

                   _lowerCAmelCase			:							Union[str, Any]      	=				model(__a					,		__a      )

                   _lowerCAmelCase			:							Dict      	=				scheduler.step(__a					,		__a					,		__a					,		generator=__a      )
                   _lowerCAmelCase			:							str      	=				output.prev_sample

            _lowerCAmelCase			:							List[Any]      	=				torch.sum(torch.abs(__a      )      )
            _lowerCAmelCase			:							List[str]      	=				torch.mean(torch.abs(__a      )      )

            assert abs(result_sum.item() - 152.3192      ) < 1E-2
            assert abs(result_mean.item() - 0.1983      ) < 1E-3







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							Union[str, Any]      	=				self.scheduler_classes[0]
            _lowerCAmelCase			:							int      	=				self.get_scheduler_config(prediction_type='v_prediction'      )
            _lowerCAmelCase			:							Any      	=				scheduler_class(**__a      )

            scheduler.set_timesteps(self.num_inference_steps      )

            _lowerCAmelCase			:							str      	=				torch.manual_seed(0      )

            _lowerCAmelCase			:							Tuple      	=				self.dummy_model()
            _lowerCAmelCase			:							Tuple      	=				self.dummy_sample_deter * scheduler.init_noise_sigma
            _lowerCAmelCase			:							Tuple      	=				sample.to(__a      )

            for i, t in enumerate(scheduler.timesteps      ):
                   _lowerCAmelCase			:							str      	=				scheduler.scale_model_input(__a					,		__a      )

                   _lowerCAmelCase			:							Union[str, Any]      	=				model(__a					,		__a      )

                   _lowerCAmelCase			:							Union[str, Any]      	=				scheduler.step(__a					,		__a					,		__a					,		generator=__a      )
                   _lowerCAmelCase			:							Optional[int]      	=				output.prev_sample

            _lowerCAmelCase			:							Optional[int]      	=				torch.sum(torch.abs(__a      )      )
            _lowerCAmelCase			:							List[Any]      	=				torch.mean(torch.abs(__a      )      )

            assert abs(result_sum.item() - 108.4439      ) < 1E-2
            assert abs(result_mean.item() - 0.1412      ) < 1E-3







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							List[Any]      	=				self.scheduler_classes[0]
            _lowerCAmelCase			:							Tuple      	=				self.get_scheduler_config()
            _lowerCAmelCase			:							Optional[int]      	=				scheduler_class(**__a      )

            scheduler.set_timesteps(self.num_inference_steps					,		device=__a      )
            _lowerCAmelCase			:							Any      	=				torch.manual_seed(0      )

            _lowerCAmelCase			:							List[str]      	=				self.dummy_model()
            _lowerCAmelCase			:							int      	=				self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
            _lowerCAmelCase			:							List[str]      	=				sample.to(__a      )

            for t in scheduler.timesteps:
                   _lowerCAmelCase			:							Optional[int]      	=				scheduler.scale_model_input(__a					,		__a      )

                   _lowerCAmelCase			:							Union[str, Any]      	=				model(__a					,		__a      )

                   _lowerCAmelCase			:							Optional[Any]      	=				scheduler.step(__a					,		__a					,		__a					,		generator=__a      )
                   _lowerCAmelCase			:							List[Any]      	=				output.prev_sample

            _lowerCAmelCase			:							int      	=				torch.sum(torch.abs(__a      )      )
            _lowerCAmelCase			:							List[Any]      	=				torch.mean(torch.abs(__a      )      )

            assert abs(result_sum.item() - 152.3192      ) < 1E-2
            assert abs(result_mean.item() - 0.1983      ) < 1E-3








"""simple docstring"""


from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


lowerCamelCase__      =      {
    "configuration_utils": ["GenerationConfig"],
    "streamers": ["TextIteratorStreamer", "TextStreamer"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__      =      [
      "Constraint",
      "ConstraintListState",
      "DisjunctiveConstraint",
      "PhrasalConstraint",
  ]
  lowerCamelCase__      =      [
      "BeamHypotheses",
      "BeamScorer",
      "BeamSearchScorer",
      "ConstrainedBeamSearchScorer",
  ]
  lowerCamelCase__      =      [
      "EpsilonLogitsWarper",
      "EtaLogitsWarper",
      "ForcedBOSTokenLogitsProcessor",
      "ForcedEOSTokenLogitsProcessor",
      "HammingDiversityLogitsProcessor",
      "InfNanRemoveLogitsProcessor",
      "LogitsProcessor",
      "LogitsProcessorList",
      "LogitsWarper",
      "MinLengthLogitsProcessor",
      "MinNewTokensLengthLogitsProcessor",
      "NoBadWordsLogitsProcessor",
      "NoRepeatNGramLogitsProcessor",
      "PrefixConstrainedLogitsProcessor",
      "RepetitionPenaltyLogitsProcessor",
      "SequenceBiasLogitsProcessor",
      "EncoderRepetitionPenaltyLogitsProcessor",
      "TemperatureLogitsWarper",
      "TopKLogitsWarper",
      "TopPLogitsWarper",
      "TypicalLogitsWarper",
      "EncoderNoRepeatNGramLogitsProcessor",
      "ExponentialDecayLengthPenalty",
      "LogitNormalization",
  ]
  lowerCamelCase__      =      [
      "MaxNewTokensCriteria",
      "MaxLengthCriteria",
      "MaxTimeCriteria",
      "StoppingCriteria",
      "StoppingCriteriaList",
      "validate_stopping_criteria",
  ]
  lowerCamelCase__      =      [
      "GenerationMixin",
      "top_k_top_p_filtering",
      "GreedySearchEncoderDecoderOutput",
      "GreedySearchDecoderOnlyOutput",
      "SampleEncoderDecoderOutput",
      "SampleDecoderOnlyOutput",
      "BeamSearchEncoderDecoderOutput",
      "BeamSearchDecoderOnlyOutput",
      "BeamSampleEncoderDecoderOutput",
      "BeamSampleDecoderOnlyOutput",
      "ContrastiveSearchEncoderDecoderOutput",
      "ContrastiveSearchDecoderOnlyOutput",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__      =      [
      "TFForcedBOSTokenLogitsProcessor",
      "TFForcedEOSTokenLogitsProcessor",
      "TFLogitsProcessor",
      "TFLogitsProcessorList",
      "TFLogitsWarper",
      "TFMinLengthLogitsProcessor",
      "TFNoBadWordsLogitsProcessor",
      "TFNoRepeatNGramLogitsProcessor",
      "TFRepetitionPenaltyLogitsProcessor",
      "TFTemperatureLogitsWarper",
      "TFTopKLogitsWarper",
      "TFTopPLogitsWarper",
      "TFForceTokensLogitsProcessor",
      "TFSuppressTokensAtBeginLogitsProcessor",
      "TFSuppressTokensLogitsProcessor",
  ]
  lowerCamelCase__      =      [
      "TFGenerationMixin",
      "tf_top_k_top_p_filtering",
      "TFGreedySearchDecoderOnlyOutput",
      "TFGreedySearchEncoderDecoderOutput",
      "TFSampleEncoderDecoderOutput",
      "TFSampleDecoderOnlyOutput",
      "TFBeamSearchEncoderDecoderOutput",
      "TFBeamSearchDecoderOnlyOutput",
      "TFBeamSampleEncoderDecoderOutput",
      "TFBeamSampleDecoderOnlyOutput",
      "TFContrastiveSearchEncoderDecoderOutput",
      "TFContrastiveSearchDecoderOnlyOutput",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__      =      [
      "FlaxForcedBOSTokenLogitsProcessor",
      "FlaxForcedEOSTokenLogitsProcessor",
      "FlaxLogitsProcessor",
      "FlaxLogitsProcessorList",
      "FlaxLogitsWarper",
      "FlaxMinLengthLogitsProcessor",
      "FlaxTemperatureLogitsWarper",
      "FlaxTopKLogitsWarper",
      "FlaxTopPLogitsWarper",
  ]
  lowerCamelCase__      =      [
      "FlaxGenerationMixin",
      "FlaxGreedySearchOutput",
      "FlaxSampleOutput",
      "FlaxBeamSearchOutput",
  ]

if TYPE_CHECKING:
  from .configuration_utils import GenerationConfig
  from .streamers import TextIteratorStreamer, TextStreamer

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
    from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
    from .logits_process import (
        EncoderNoRepeatNGramLogitsProcessor,
        EncoderRepetitionPenaltyLogitsProcessor,
        EpsilonLogitsWarper,
        EtaLogitsWarper,
        ExponentialDecayLengthPenalty,
        ForcedBOSTokenLogitsProcessor,
        ForcedEOSTokenLogitsProcessor,
        HammingDiversityLogitsProcessor,
        InfNanRemoveLogitsProcessor,
        LogitNormalization,
        LogitsProcessor,
        LogitsProcessorList,
        LogitsWarper,
        MinLengthLogitsProcessor,
        MinNewTokensLengthLogitsProcessor,
        NoBadWordsLogitsProcessor,
        NoRepeatNGramLogitsProcessor,
        PrefixConstrainedLogitsProcessor,
        RepetitionPenaltyLogitsProcessor,
        SequenceBiasLogitsProcessor,
        TemperatureLogitsWarper,
        TopKLogitsWarper,
        TopPLogitsWarper,
        TypicalLogitsWarper,
    )
    from .stopping_criteria import (
        MaxLengthCriteria,
        MaxNewTokensCriteria,
        MaxTimeCriteria,
        StoppingCriteria,
        StoppingCriteriaList,
        validate_stopping_criteria,
    )
    from .utils import (
        BeamSampleDecoderOnlyOutput,
        BeamSampleEncoderDecoderOutput,
        BeamSearchDecoderOnlyOutput,
        BeamSearchEncoderDecoderOutput,
        ContrastiveSearchDecoderOnlyOutput,
        ContrastiveSearchEncoderDecoderOutput,
        GenerationMixin,
        GreedySearchDecoderOnlyOutput,
        GreedySearchEncoderDecoderOutput,
        SampleDecoderOnlyOutput,
        SampleEncoderDecoderOutput,
        top_k_top_p_filtering,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tf_logits_process import (
        TFForcedBOSTokenLogitsProcessor,
        TFForcedEOSTokenLogitsProcessor,
        TFForceTokensLogitsProcessor,
        TFLogitsProcessor,
        TFLogitsProcessorList,
        TFLogitsWarper,
        TFMinLengthLogitsProcessor,
        TFNoBadWordsLogitsProcessor,
        TFNoRepeatNGramLogitsProcessor,
        TFRepetitionPenaltyLogitsProcessor,
        TFSuppressTokensAtBeginLogitsProcessor,
        TFSuppressTokensLogitsProcessor,
        TFTemperatureLogitsWarper,
        TFTopKLogitsWarper,
        TFTopPLogitsWarper,
    )
    from .tf_utils import (
        TFBeamSampleDecoderOnlyOutput,
        TFBeamSampleEncoderDecoderOutput,
        TFBeamSearchDecoderOnlyOutput,
        TFBeamSearchEncoderDecoderOutput,
        TFContrastiveSearchDecoderOnlyOutput,
        TFContrastiveSearchEncoderDecoderOutput,
        TFGenerationMixin,
        TFGreedySearchDecoderOnlyOutput,
        TFGreedySearchEncoderDecoderOutput,
        TFSampleDecoderOnlyOutput,
        TFSampleEncoderDecoderOutput,
        tf_top_k_top_p_filtering,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .flax_logits_process import (
        FlaxForcedBOSTokenLogitsProcessor,
        FlaxForcedEOSTokenLogitsProcessor,
        FlaxLogitsProcessor,
        FlaxLogitsProcessorList,
        FlaxLogitsWarper,
        FlaxMinLengthLogitsProcessor,
        FlaxTemperatureLogitsWarper,
        FlaxTopKLogitsWarper,
        FlaxTopPLogitsWarper,
    )
    from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
  import sys

  lowerCamelCase__      =      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
