




'''simple docstring'''



import json
import os
import pickle
import subprocess
from hashlib import mda
from pathlib import Path
from tempfile import gettempdir
from textwrap import dedent
from types import CodeType, FunctionType
from unittest import TestCase
from unittest.mock import patch

import pytest
from multiprocess import Pool
from packaging import version

import datasets
from datasets.fingerprint import Hasher, fingerprint_transform
from datasets.table import InMemoryTable

from .utils import (
    require_regex,
    require_spacy,
    require_spacy_model,
    require_tiktoken,
    require_torch,
    require_transformers,
)







class 		lowercase_    :


   """simple docstring"""

   def __init__(    self					:	Optional[Any]  ,      __lowerCamelCase					:	Tuple    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      foo



   def __call__(    self					:	List[str]    ):

      """simple docstring"""
      return self.foo







class 		lowercase_    ( datasets.Dataset   ):


   """simple docstring"""

   @fingerprint_transform(inplace=_snake_case    )
   def 							lowerCAmelCase_			(    self					:	Optional[int]  ,      __lowerCamelCase					:	List[Any]  ,      *__lowerCamelCase					:	Dict  ,      **__lowerCamelCase					:	List[str]    ):

      """simple docstring"""
      return DatasetChild(self.data  ,      fingerprint=_snake_case    )



   @fingerprint_transform(inplace=_snake_case    )
   def 							lowerCAmelCase_			(    self					:	Union[str, Any]  ,      __lowerCamelCase					:	Tuple  ,      *__lowerCamelCase					:	Tuple  ,      **__lowerCamelCase					:	Optional[Any]    ):

      """simple docstring"""
      return DatasetChild(self.data  ,      fingerprint=_snake_case    )







class 		lowercase_    :


   """simple docstring"""

   def __init__(    self					:	List[Any]  ,      __lowerCamelCase					:	List[Any]    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      callable



   def __call__(    self					:	int  ,      *__lowerCamelCase					:	int  ,      **__lowerCamelCase					:	str    ):

      """simple docstring"""






      if self.callable is not None:
         return self.callable(*_snake_case  ,      **_snake_case    )



   def __getstate__(    self					:	Union[str, Any]    ):

      """simple docstring"""
      raise pickle.PicklingError()







class 		lowercase_    ( a__   ):


   """simple docstring"""

   @require_transformers
   @pytest.mark.integration
   def 							lowerCAmelCase_			(    self					:	Tuple    ):

      """simple docstring"""
      from transformers import AutoTokenizer

      def encode(__lowerCamelCase					:	str    ):
         return tokenizer(_snake_case    )

      # TODO: add hash consistency tests across sessions
      _SCREAMING_SNAKE_CASE       =      AutoTokenizer.from_pretrained("bert-base-uncased"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(lambda __lowerCamelCase    : tokenizer(_snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      AutoTokenizer.from_pretrained("bert-base-cased"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(lambda __lowerCamelCase    : tokenizer(_snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      AutoTokenizer.from_pretrained("bert-base-uncased"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(lambda __lowerCamelCase    : tokenizer(_snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   @require_transformers
   @pytest.mark.integration
   def 							lowerCAmelCase_			(    self					:	Dict    ):

      """simple docstring"""
      from transformers import AutoTokenizer

      _SCREAMING_SNAKE_CASE       =      AutoTokenizer.from_pretrained("gpt2"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      tokenizer("Hello world !"    )  # call once to change the tokenizer's cache
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )



   @require_regex
   def 							lowerCAmelCase_			(    self					:	Dict    ):

      """simple docstring"""
      import regex

      _SCREAMING_SNAKE_CASE       =      regex.Regex("foo"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      regex.Regex("bar"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      regex.Regex("foo"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )







class 		lowercase_    ( a__   ):


   """simple docstring"""

   def 							lowerCAmelCase_			(    self					:	Union[str, Any]    ):

      """simple docstring"""






      def func():
         return foo

      _SCREAMING_SNAKE_CASE       =      [0]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      [1]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      [0]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Any    ):

      """simple docstring"""






      def func():
         pass

      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()

      def func():
         pass

      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Tuple    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(Foo([0]    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(Foo([1]    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(Foo([0]    )    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Dict    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(Foo([0]    ).__call__    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(Foo([1]    ).__call__    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(Foo([0]    ).__call__    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Optional[Any]    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      (
          "co_argcount",
          "co_kwonlyargcount",
          "co_nlocals",
          "co_stacksize",
          "co_flags",
          "co_code",
          "co_consts",
          "co_names",
          "co_varnames",
          "co_filename",
          "co_name",
          "co_firstlineno",
          "co_lnotab",
          "co_freevars",
          "co_cellvars",
      )

      def _create_code_pyaa(*__lowerCamelCase					:	Dict    ):
         if hasattr(args[-3]  ,      "encode"    ):
            _SCREAMING_SNAKE_CASE       =      list(_snake_case    )
            _SCREAMING_SNAKE_CASE       =      args[-3].encode()  # co_lnotab
            _SCREAMING_SNAKE_CASE       =      args[-1_0].encode()  # co_code
         if hasattr(_snake_case  ,      "co_posonlyargcount"    ):
            if len(_snake_case    ) == 1_6:
               return CodeType(*_snake_case    )
            elif len(_snake_case    ) == 1_5:
               return CodeType(args[0]  ,      0  ,      *args[1:]    )
            return CodeType(args[0]  ,      0  ,      0  ,      *args[1:]    )
         elif hasattr(_snake_case  ,      "co_kwonlyargcount"    ):
            if len(_snake_case    ) == 1_6:
               return CodeType(args[0]  ,      *args[2:]    )
            elif len(_snake_case    ) == 1_5:
               return CodeType(*_snake_case    )
            return CodeType(args[0]  ,      0  ,      *args[1:]    )
         if len(_snake_case    ) == 1_6:
            return CodeType(args[0]  ,      *args[3:]    )
         elif len(_snake_case    ) == 1_5:
            return CodeType(args[0]  ,      *args[2:]    )
         return CodeType(*_snake_case    )

      def create_ipython_func(__lowerCamelCase					:	Tuple  ,      __lowerCamelCase					:	List[str]    ):
         def func():
            return returned_obj

         _SCREAMING_SNAKE_CASE       =      func.__code__
         # Use _create_code from dill in order to make it work for different python versions
         if datasets.config.PY_VERSION < version.parse("3.8"    ):
            _SCREAMING_SNAKE_CASE       =      _create_code_pyaa(
                *[getattr(_snake_case  ,      _snake_case    ) if k != "co_filename" else co_filename for k in code_args_pyaa]    )
         else:
            _SCREAMING_SNAKE_CASE       =      code.replace(co_filename=_snake_case    )
         return FunctionType(_snake_case  ,      func.__globals__  ,      func.__name__  ,      func.__defaults__  ,      func.__closure__    )

      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      "<ipython-input-2-e0383a102aae>", [0]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(create_ipython_func(_snake_case  ,      _snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      "<ipython-input-2-e0383a102aae>", [1]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(create_ipython_func(_snake_case  ,      _snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      "<ipython-input-5-713f6613acf3>", [0]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(create_ipython_func(_snake_case  ,      _snake_case    )    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )

      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      os.path.join(gettempdir()  ,      "ipykernel_12345"  ,      "321456789.py"    ), [0]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(create_ipython_func(_snake_case  ,      _snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      os.path.join(gettempdir()  ,      "ipykernel_12345"  ,      "321456789.py"    ), [1]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(create_ipython_func(_snake_case  ,      _snake_case    )    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      os.path.join(gettempdir()  ,      "ipykernel_12345"  ,      "654123987.py"    ), [0]
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(create_ipython_func(_snake_case  ,      _snake_case    )    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Tuple    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE       =      [0], [1]

      def func():
         return foo, bar

      _SCREAMING_SNAKE_CASE       =      "__main__"

      def globalvars_mocka_side_effect(__lowerCamelCase					:	Dict  ,      *__lowerCamelCase					:	Any  ,      **__lowerCamelCase					:	Dict    ):
         return {"foo": foo, "bar": bar}

      def globalvars_mocka_side_effect(__lowerCamelCase					:	Optional[int]  ,      *__lowerCamelCase					:	Optional[int]  ,      **__lowerCamelCase					:	Optional[int]    ):
         return {"bar": bar, "foo": foo}

      with patch("dill.detect.globalvars"  ,      side_effect=_snake_case    ) as globalvars_mocka:
         _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
         self.assertGreater(globalvars_mocka.call_count  ,      0    )
      with patch("dill.detect.globalvars"  ,      side_effect=_snake_case    ) as globalvars_mocka:
         _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
         self.assertGreater(globalvars_mocka.call_count  ,      0    )
      self.assertEqual(_snake_case  ,      _snake_case    )







class 		lowercase_    ( a__   ):


   """simple docstring"""

   def 							lowerCAmelCase_			(    self					:	Tuple    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      Hasher.hash("hello"    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash("hello"    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash("there"    )
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Any    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      Hasher.hash(Foo("hello"    )    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash(Foo("hello"    )    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash(Foo("there"    )    )
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Union[str, Any]    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      Hasher()
      for x in ["hello", Foo("hello"    )]:
         hasher.update(_snake_case    )
      _SCREAMING_SNAKE_CASE       =      hasher.hexdigest()
      _SCREAMING_SNAKE_CASE       =      Hasher()
      for x in ["hello", Foo("hello"    )]:
         hasher.update(_snake_case    )
      _SCREAMING_SNAKE_CASE       =      hasher.hexdigest()
      _SCREAMING_SNAKE_CASE       =      Hasher()
      for x in ["there", Foo("there"    )]:
         hasher.update(_snake_case    )
      _SCREAMING_SNAKE_CASE       =      hasher.hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   def 							lowerCAmelCase_			(    self					:	Dict    ):

      """simple docstring"""






      with self.assertRaises(pickle.PicklingError    ):
         Hasher.hash(UnpicklableCallable(Foo("hello"    )    )    )



   def 							lowerCAmelCase_			(    self					:	List[str]    ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE       =      "abc"
      _SCREAMING_SNAKE_CASE       =      [string, string]  # two strings have the same ids
      _SCREAMING_SNAKE_CASE       =      [string, string]
      _SCREAMING_SNAKE_CASE       =      json.loads(F"""[\"{string}\", \"{string}\"]"""    )  # two strings have different ids
      self.assertIs(obja[0]  ,      _snake_case    )
      self.assertIs(obja[0]  ,      obja[1]    )
      self.assertIs(obja[0]  ,      _snake_case    )
      self.assertIs(obja[0]  ,      obja[1]    )
      self.assertIsNot(obja[0]  ,      _snake_case    )
      self.assertIsNot(obja[0]  ,      obja[1]    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash(_snake_case    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash(_snake_case    )
      _SCREAMING_SNAKE_CASE       =      Hasher.hash(_snake_case    )
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertEqual(_snake_case  ,      _snake_case    )



   @require_tiktoken
   def 							lowerCAmelCase_			(    self					:	List[Any]    ):

      """simple docstring"""
      import tiktoken

      _SCREAMING_SNAKE_CASE       =      tiktoken.get_encoding("gpt2"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      tiktoken.get_encoding("r50k_base"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      tiktoken.get_encoding("gpt2"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   @require_torch
   def 							lowerCAmelCase_			(    self					:	Tuple    ):

      """simple docstring"""
      import torch

      _SCREAMING_SNAKE_CASE       =      torch.tensor([1.0]    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      torch.tensor([2.0]    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      torch.tensor([1.0]    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )



   @require_spacy
   @require_spacy_model("en_core_web_sm"    )
   @require_spacy_model("fr_core_news_sm"    )
   @pytest.mark.integration
   def 							lowerCAmelCase_			(    self					:	int    ):

      """simple docstring"""
      import spacy

      _SCREAMING_SNAKE_CASE       =      spacy.load("en_core_web_sm"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      spacy.load("fr_core_news_sm"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      _SCREAMING_SNAKE_CASE       =      spacy.load("en_core_web_sm"    )
      _SCREAMING_SNAKE_CASE       =      mda(datasets.utils.py_utils.dumps(_snake_case    )    ).hexdigest()
      self.assertEqual(_snake_case  ,      _snake_case    )
      self.assertNotEqual(_snake_case  ,      _snake_case    )


@pytest.mark.integration
def  SCREAMING_SNAKE_CASE_ (							__A				: Path  ) -> Any:
   _SCREAMING_SNAKE_CASE       =      tmp_path / "dir1"
   _SCREAMING_SNAKE_CASE       =      tmp_path / "dir2"
   dira.mkdir()
   dira.mkdir()
   _SCREAMING_SNAKE_CASE       =      "script.py"
   _SCREAMING_SNAKE_CASE       =      dedent(
       "\n    from datasets.fingerprint import Hasher\n    def foo():\n        pass\n    print(Hasher.hash(foo))\n    "  )
   _SCREAMING_SNAKE_CASE       =      dira / script_filename
   _SCREAMING_SNAKE_CASE       =      dira / script_filename
   with script_patha.open("w"  ) as f:
      f.write(_UpperCAmelCase  )
   with script_patha.open("w"  ) as f:
      f.write(_UpperCAmelCase  )
   _SCREAMING_SNAKE_CASE       =      subprocess.check_output(["python", str(_UpperCAmelCase  )]  )
   _SCREAMING_SNAKE_CASE       =      subprocess.check_output(["python", str(_UpperCAmelCase  )]  )
   assert fingerprinta == fingerprinta
def  SCREAMING_SNAKE_CASE_ (							) -> List[str]:
   _SCREAMING_SNAKE_CASE       =      {"a": [0, 1, 2]}
   _SCREAMING_SNAKE_CASE       =      DatasetChild(InMemoryTable.from_pydict(_UpperCAmelCase  )  )
   _SCREAMING_SNAKE_CASE       =      dataset.funca()._fingerprint
   assert expected_fingerprint == dataset.funca()._fingerprint
   assert expected_fingerprint != dataset.funca()._fingerprint

   with Pool(2  ) as p:
      assert expected_fingerprint == p.apply_async(dataset.funca  ).get()._fingerprint
      assert expected_fingerprint != p.apply_async(dataset.funca  ).get()._fingerprint
def  SCREAMING_SNAKE_CASE_ (							) -> Optional[Any]:
   _SCREAMING_SNAKE_CASE       =      {"a": [0, 1, 2]}







   class 		lowercase_    ( datasets.Dataset   ):


      """simple docstring"""

      @fingerprint_transform(inplace=_snake_case    )
      def 							lowerCAmelCase_			(    self					:	Any  ,      __lowerCamelCase					:	List[str]    ):

         """simple docstring"""
         return DummyDatasetChild(self.data  ,      fingerprint=_snake_case    )


   _SCREAMING_SNAKE_CASE       =      DummyDatasetChild(InMemoryTable.from_pydict(_UpperCAmelCase  )  ).func()







   class 		lowercase_    ( datasets.Dataset   ):


      """simple docstring"""

      @fingerprint_transform(inplace=_snake_case  ,      version="1.0.0"    )
      def 							lowerCAmelCase_			(    self					:	Tuple  ,      __lowerCamelCase					:	Tuple    ):

         """simple docstring"""
         return DummyDatasetChild(self.data  ,      fingerprint=_snake_case    )


   _SCREAMING_SNAKE_CASE       =      DummyDatasetChild(InMemoryTable.from_pydict(_UpperCAmelCase  )  ).func()







   class 		lowercase_    ( datasets.Dataset   ):


      """simple docstring"""

      @fingerprint_transform(inplace=_snake_case  ,      version="2.0.0"    )
      def 							lowerCAmelCase_			(    self					:	Optional[int]  ,      __lowerCamelCase					:	Union[str, Any]    ):

         """simple docstring"""
         return DummyDatasetChild(self.data  ,      fingerprint=_snake_case    )


   _SCREAMING_SNAKE_CASE       =      DummyDatasetChild(InMemoryTable.from_pydict(_UpperCAmelCase  )  ).func()

   assert len({fingeprint_no_version, fingeprint_a, fingeprint_a}  ) == 3






def  SCREAMING_SNAKE_CASE_ (							) -> List[str]:
   # AttributeError: module 'dill._dill' has no attribute 'stack'
   _SCREAMING_SNAKE_CASE       =      Hasher()
   hasher.update(lambda __A  : x  )








"""simple docstring"""
import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class 	a				(					a__   ):
 snake_case__    	=    (HeunDiscreteScheduler,)
 snake_case__    	=    1_0



 def 					UpperCamelCase__   (			self							,    **_snake_case			):







  """simple docstring"""






  lowerCAmelCase  							=      {
      'num_train_timesteps': 11_00,
      'beta_start': 0.0_001,
      'beta_end': 0.02,
      'beta_schedule': 'linear',
  }

  config.update(**_snake_case			)
  return config



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  for timesteps in [10, 50, 1_00, 10_00]:
   self.check_over_configs(num_train_timesteps=_snake_case			)



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  for beta_start, beta_end in zip([0.00_001, 0.0_001, 0.001]							,    [0.0_002, 0.002, 0.02]			):
   self.check_over_configs(beta_start=_snake_case							,    beta_end=_snake_case			)



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






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



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  for clip_sample_range in [1.0, 2.0, 3.0]:
   self.check_over_configs(clip_sample_range=_snake_case							,    clip_sample=_snake_case			)



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






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



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  lowerCAmelCase  							=      self.scheduler_classes[0]
  lowerCAmelCase  							=      self.get_scheduler_config()
  lowerCAmelCase  							=      scheduler_class(**_snake_case			)

  scheduler.set_timesteps(self.num_inference_steps			)

  lowerCAmelCase  							=      self.dummy_model()
  lowerCAmelCase  							=      self.dummy_sample_deter * scheduler.init_noise_sigma
  lowerCAmelCase  							=      sample.to(_snake_case			)

  for i, t in enumerate(scheduler.timesteps			):
   lowerCAmelCase  							=      scheduler.scale_model_input(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      model(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      scheduler.step(_snake_case							,    _snake_case							,    _snake_case			)
   lowerCAmelCase  							=      output.prev_sample

  lowerCAmelCase  							=      torch.sum(torch.abs(_snake_case			)			)
  lowerCAmelCase  							=      torch.mean(torch.abs(_snake_case			)			)

  if torch_device in ["cpu", "mps"]:
   assert abs(result_sum.item() - 0.1_233			) < 1E-2
   assert abs(result_mean.item() - 0.0_002			) < 1E-3
  else:
   # CUDA
   assert abs(result_sum.item() - 0.1_233			) < 1E-2
   assert abs(result_mean.item() - 0.0_002			) < 1E-3



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  lowerCAmelCase  							=      self.scheduler_classes[0]
  lowerCAmelCase  							=      self.get_scheduler_config(prediction_type='v_prediction'			)
  lowerCAmelCase  							=      scheduler_class(**_snake_case			)

  scheduler.set_timesteps(self.num_inference_steps			)

  lowerCAmelCase  							=      self.dummy_model()
  lowerCAmelCase  							=      self.dummy_sample_deter * scheduler.init_noise_sigma
  lowerCAmelCase  							=      sample.to(_snake_case			)

  for i, t in enumerate(scheduler.timesteps			):
   lowerCAmelCase  							=      scheduler.scale_model_input(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      model(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      scheduler.step(_snake_case							,    _snake_case							,    _snake_case			)
   lowerCAmelCase  							=      output.prev_sample

  lowerCAmelCase  							=      torch.sum(torch.abs(_snake_case			)			)
  lowerCAmelCase  							=      torch.mean(torch.abs(_snake_case			)			)

  if torch_device in ["cpu", "mps"]:
   assert abs(result_sum.item() - 4.69_34E-07			) < 1E-2
   assert abs(result_mean.item() - 6.11_12E-10			) < 1E-3
  else:
   # CUDA
   assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72E-07			) < 1E-2
   assert abs(result_mean.item() - 0.0_002			) < 1E-3



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  lowerCAmelCase  							=      self.scheduler_classes[0]
  lowerCAmelCase  							=      self.get_scheduler_config()
  lowerCAmelCase  							=      scheduler_class(**_snake_case			)

  scheduler.set_timesteps(self.num_inference_steps							,    device=_snake_case			)

  lowerCAmelCase  							=      self.dummy_model()
  lowerCAmelCase  							=      self.dummy_sample_deter.to(_snake_case			) * scheduler.init_noise_sigma

  for t in scheduler.timesteps:
   lowerCAmelCase  							=      scheduler.scale_model_input(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      model(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      scheduler.step(_snake_case							,    _snake_case							,    _snake_case			)
   lowerCAmelCase  							=      output.prev_sample

  lowerCAmelCase  							=      torch.sum(torch.abs(_snake_case			)			)
  lowerCAmelCase  							=      torch.mean(torch.abs(_snake_case			)			)

  if str(_snake_case			).startswith('cpu'			):
   # The following sum varies between 148 and 156 on mps. Why?
   assert abs(result_sum.item() - 0.1_233			) < 1E-2
   assert abs(result_mean.item() - 0.0_002			) < 1E-3
  elif str(_snake_case			).startswith('mps'			):
   # Larger tolerance on mps
   assert abs(result_mean.item() - 0.0_002			) < 1E-2
  else:
   # CUDA
   assert abs(result_sum.item() - 0.1_233			) < 1E-2
   assert abs(result_mean.item() - 0.0_002			) < 1E-3



 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






  lowerCAmelCase  							=      self.scheduler_classes[0]
  lowerCAmelCase  							=      self.get_scheduler_config()
  lowerCAmelCase  							=      scheduler_class(**_snake_case							,    use_karras_sigmas=_snake_case			)

  scheduler.set_timesteps(self.num_inference_steps							,    device=_snake_case			)

  lowerCAmelCase  							=      self.dummy_model()
  lowerCAmelCase  							=      self.dummy_sample_deter.to(_snake_case			) * scheduler.init_noise_sigma
  lowerCAmelCase  							=      sample.to(_snake_case			)

  for t in scheduler.timesteps:
   lowerCAmelCase  							=      scheduler.scale_model_input(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      model(_snake_case							,    _snake_case			)

   lowerCAmelCase  							=      scheduler.step(_snake_case							,    _snake_case							,    _snake_case			)
   lowerCAmelCase  							=      output.prev_sample

  lowerCAmelCase  							=      torch.sum(torch.abs(_snake_case			)			)
  lowerCAmelCase  							=      torch.mean(torch.abs(_snake_case			)			)

  assert abs(result_sum.item() - 0.00_015			) < 1E-2
  assert abs(result_mean.item() - 1.98_69_55_45_35_03_46_95E-07			) < 1E-2





