code: 
  '''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  lowerCAmelCase__   :
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def __init__(self  : int	, lowerCAmelCase_  : Optional[Any]    )   ->       Tuple:
          snake_case_						       =							foo
  
  
    def __call__(self  : List[Any]    )   ->       Optional[int]:
          return self.foo
  
  
  
  class  lowerCAmelCase__   (       datasets.Dataset       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    @fingerprint_transform(inplace=lowerCAmelCase_    )
    def 		_lowerCamelCase	(self  : Tuple	, lowerCAmelCase_  : Any	, *lowerCAmelCase_  : Union[str, Any]	, **lowerCAmelCase_  : Any    )   ->       str:
          return DatasetChild(self.data	, fingerprint=lowerCAmelCase_    )
  
  
    @fingerprint_transform(inplace=lowerCAmelCase_    )
    def 		_lowerCamelCase	(self  : str	, lowerCAmelCase_  : Optional[Any]	, *lowerCAmelCase_  : List[Any]	, **lowerCAmelCase_  : Any    )   ->       List[str]:
          return DatasetChild(self.data	, fingerprint=lowerCAmelCase_    )
  
  
  
  class  lowerCAmelCase__   :
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def __init__(self  : int	, lowerCAmelCase_  : List[str]    )   ->       Dict:
          snake_case_						       =							callable
  
    def __call__(self  : List[str]	, *lowerCAmelCase_  : Any	, **lowerCAmelCase_  : Tuple    )   ->       Dict:
          if self.callable is not None:
                return self.callable(*lowerCAmelCase_	, **lowerCAmelCase_    )
  
  
    def __getstate__(self  : List[str]    )   ->       Tuple:
          raise pickle.PicklingError()
  
  
  
  class  lowerCAmelCase__   (       UpperCamelCase       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    @require_transformers
    @pytest.mark.integration
    def 		_lowerCamelCase	(self  : str    )   ->       str:
          from transformers import AutoTokenizer
  
          def encode(lowerCAmelCase_  : str    ):
                return tokenizer(lowerCAmelCase_    )
  
          # TODO: add hash consistency tests across sessions
          snake_case_						       =							AutoTokenizer.from_pretrained("bert-base-uncased"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lambda lowerCAmelCase_    : tokenizer(lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							AutoTokenizer.from_pretrained("bert-base-cased"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lambda lowerCAmelCase_    : tokenizer(lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							AutoTokenizer.from_pretrained("bert-base-uncased"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lambda lowerCAmelCase_    : tokenizer(lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    @require_transformers
    @pytest.mark.integration
    def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Optional[int]:
          from transformers import AutoTokenizer
  
          snake_case_						       =							AutoTokenizer.from_pretrained("gpt2"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          tokenizer("Hello world !"    )  # call once to change the tokenizer's cache
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
  
    @require_regex
    def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       Optional[int]:
          import regex
  
          snake_case_						       =							regex.Regex("foo"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							regex.Regex("bar"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							regex.Regex("foo"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
  
  
  class  lowerCAmelCase__   (       UpperCamelCase       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def 		_lowerCamelCase	(self  : Dict    )   ->       Optional[Any]:
          def func():
                return foo
  
          snake_case_						       =							[0]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							[1]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							[0]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : str    )   ->       Tuple:
          def func():
                pass
  
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
  
          def func():
                pass
  
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : List[Any]    )   ->       Union[str, Any]:
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(Foo([0]    )    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(Foo([1]    )    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(Foo([0]    )    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       int:
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(Foo([0]    ).__call__    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(Foo([1]    ).__call__    )    ).hexdigest()
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(Foo([0]    ).__call__    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : int    )   ->       Optional[Any]:
          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_  : Any    ):
                if hasattr(args[-3]	, "encode"    ):
                      snake_case_						       =							list(lowerCAmelCase_    )
                      snake_case_						       =							args[-3].encode()  # co_lnotab
                      snake_case_						       =							args[-10].encode()  # co_code
                if hasattr(lowerCAmelCase_	, "co_posonlyargcount"    ):
                      if len(lowerCAmelCase_    ) == 16:
                            return CodeType(*lowerCAmelCase_    )
                      elif len(lowerCAmelCase_    ) == 15:
                            return CodeType(args[0]	, 0	, *args[1:]    )
                      return CodeType(args[0]	, 0	, 0	, *args[1:]    )
                elif hasattr(lowerCAmelCase_	, "co_kwonlyargcount"    ):
                      if len(lowerCAmelCase_    ) == 16:
                            return CodeType(args[0]	, *args[2:]    )
                      elif len(lowerCAmelCase_    ) == 15:
                            return CodeType(*lowerCAmelCase_    )
                      return CodeType(args[0]	, 0	, *args[1:]    )
                if len(lowerCAmelCase_    ) == 16:
                      return CodeType(args[0]	, *args[3:]    )
                elif len(lowerCAmelCase_    ) == 15:
                      return CodeType(args[0]	, *args[2:]    )
                return CodeType(*lowerCAmelCase_    )
  
          def create_ipython_func(lowerCAmelCase_  : Tuple	, lowerCAmelCase_  : Any    ):
                def func():
                      return returned_obj
  
                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"    ):
                      snake_case_						       =							_create_code_pyaa(
                          *[getattr(lowerCAmelCase_	, lowerCAmelCase_    ) if k != "co_filename" else co_filename for k in code_args_pyaa]    )
                else:
                      snake_case_						       =							code.replace(co_filename=lowerCAmelCase_    )
                return FunctionType(lowerCAmelCase_	, func.__globals__	, func.__name__	, func.__defaults__	, func.__closure__    )
  
          snake_case_       ,	snake_case_						       =							"<ipython-input-2-e0383a102aae>", [0]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(create_ipython_func(lowerCAmelCase_	, lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_       ,	snake_case_						       =							"<ipython-input-2-e0383a102aae>", [1]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(create_ipython_func(lowerCAmelCase_	, lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_       ,	snake_case_						       =							"<ipython-input-5-713f6613acf3>", [0]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(create_ipython_func(lowerCAmelCase_	, lowerCAmelCase_    )    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
          snake_case_       ,	snake_case_						       =							os.path.join(gettempdir()	, "ipykernel_12345"	, "321456789.py"    ), [0]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(create_ipython_func(lowerCAmelCase_	, lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_       ,	snake_case_						       =							os.path.join(gettempdir()	, "ipykernel_12345"	, "321456789.py"    ), [1]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(create_ipython_func(lowerCAmelCase_	, lowerCAmelCase_    )    )    ).hexdigest()
          snake_case_       ,	snake_case_						       =							os.path.join(gettempdir()	, "ipykernel_12345"	, "654123987.py"    ), [0]
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(create_ipython_func(lowerCAmelCase_	, lowerCAmelCase_    )    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
  
    def 		_lowerCamelCase	(self  : int    )   ->       List[Any]:
          snake_case_       ,	snake_case_						       =							[0], [1]
  
          def func():
                return foo, bar
  
          snake_case_						       =							"__main__"
  
          def globalvars_mocka_side_effect(lowerCAmelCase_  : str	, *lowerCAmelCase_  : Tuple	, **lowerCAmelCase_  : List[Any]    ):
                return {"foo": foo, "bar": bar}
  
          def globalvars_mocka_side_effect(lowerCAmelCase_  : List[str]	, *lowerCAmelCase_  : List[str]	, **lowerCAmelCase_  : Optional[Any]    ):
                return {"bar": bar, "foo": foo}
  
          with patch("dill.detect.globalvars"	, side_effect=lowerCAmelCase_    ) as globalvars_mocka:
                snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
                self.assertGreater(globalvars_mocka.call_count	, 0    )
          with patch("dill.detect.globalvars"	, side_effect=lowerCAmelCase_    ) as globalvars_mocka:
                snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
                self.assertGreater(globalvars_mocka.call_count	, 0    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
  
  
  class  lowerCAmelCase__   (       UpperCamelCase       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       Tuple:
          snake_case_						       =							Hasher.hash("hello"    )
          snake_case_						       =							Hasher.hash("hello"    )
          snake_case_						       =							Hasher.hash("there"    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : Any    )   ->       int:
          snake_case_						       =							Hasher.hash(Foo("hello"    )    )
          snake_case_						       =							Hasher.hash(Foo("hello"    )    )
          snake_case_						       =							Hasher.hash(Foo("there"    )    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : List[Any]    )   ->       List[Any]:
          snake_case_						       =							Hasher()
          for x in ["hello", Foo("hello"    )]:
                hasher.update(lowerCAmelCase_    )
          snake_case_						       =							hasher.hexdigest()
          snake_case_						       =							Hasher()
          for x in ["hello", Foo("hello"    )]:
                hasher.update(lowerCAmelCase_    )
          snake_case_						       =							hasher.hexdigest()
          snake_case_						       =							Hasher()
          for x in ["there", Foo("there"    )]:
                hasher.update(lowerCAmelCase_    )
          snake_case_						       =							hasher.hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : str    )   ->       List[Any]:
          with self.assertRaises(pickle.PicklingError    ):
                Hasher.hash(UnpicklableCallable(Foo("hello"    )    )    )
  
    def 		_lowerCamelCase	(self  : Dict    )   ->       Optional[int]:
          snake_case_						       =							"abc"
          snake_case_						       =							[string, string]  # two strings have the same ids
          snake_case_						       =							[string, string]
          snake_case_						       =							json.loads(f'''["{string}", "{string}"]'''    )  # two strings have different ids
          self.assertIs(obja[0]	, lowerCAmelCase_    )
          self.assertIs(obja[0]	, obja[1]    )
          self.assertIs(obja[0]	, lowerCAmelCase_    )
          self.assertIs(obja[0]	, obja[1]    )
          self.assertIsNot(obja[0]	, lowerCAmelCase_    )
          self.assertIsNot(obja[0]	, obja[1]    )
          snake_case_						       =							Hasher.hash(lowerCAmelCase_    )
          snake_case_						       =							Hasher.hash(lowerCAmelCase_    )
          snake_case_						       =							Hasher.hash(lowerCAmelCase_    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    @require_tiktoken
    def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       Union[str, Any]:
          import tiktoken
  
          snake_case_						       =							tiktoken.get_encoding("gpt2"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							tiktoken.get_encoding("r50k_base"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							tiktoken.get_encoding("gpt2"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
    @require_torch
    def 		_lowerCamelCase	(self  : Optional[int]    )   ->       str:
          import torch
  
          snake_case_						       =							torch.tensor([1.0]    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							torch.tensor([2.0]    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							torch.tensor([1.0]    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
  
    @require_spacy
    @require_spacy_model("en_core_web_sm"    )
    @require_spacy_model("fr_core_news_sm"    )
    @pytest.mark.integration
    def 		_lowerCamelCase	(self  : int    )   ->       Union[str, Any]:
          import spacy
  
          snake_case_						       =							spacy.load("en_core_web_sm"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							spacy.load("fr_core_news_sm"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          snake_case_						       =							spacy.load("en_core_web_sm"    )
          snake_case_						       =							mda(datasets.utils.py_utils.dumps(lowerCAmelCase_    )    ).hexdigest()
          self.assertEqual(lowerCAmelCase_	, lowerCAmelCase_    )
          self.assertNotEqual(lowerCAmelCase_	, lowerCAmelCase_    )
  
  @pytest.mark.integration
  def   snake_case_       (							_a     :							Path			):
        snake_case_						       =							tmp_path / "dir1"
        snake_case_						       =							tmp_path / "dir2"
        dira.mkdir()
        dira.mkdir()
        snake_case_						       =							"script.py"
        snake_case_						       =							dedent(
            "\n    from datasets.fingerprint import Hasher\n    def foo():\n        pass\n    print(Hasher.hash(foo))\n    "			)
        snake_case_						       =							dira / script_filename
        snake_case_						       =							dira / script_filename
        with script_patha.open("w"			) as f:
              f.write(_a			)
        with script_patha.open("w"			) as f:
              f.write(_a			)
        snake_case_						       =							subprocess.check_output(["python", str(_a			)]			)
        snake_case_						       =							subprocess.check_output(["python", str(_a			)]			)
        assert fingerprinta == fingerprinta
  def   snake_case_       (							):
        snake_case_						       =							{"a": [0, 1, 2]}
        snake_case_						       =							DatasetChild(InMemoryTable.from_pydict(_a			)			)
        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   snake_case_       (							):
        snake_case_						       =							{"a": [0, 1, 2]}
  
  
  
        class  lowerCAmelCase__   (       datasets.Dataset       ):
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          @fingerprint_transform(inplace=lowerCAmelCase_    )
          def 		_lowerCamelCase	(self  : int	, lowerCAmelCase_  : str    )   ->       List[Any]:
                return DummyDatasetChild(self.data	, fingerprint=lowerCAmelCase_    )
  
        snake_case_						       =							DummyDatasetChild(InMemoryTable.from_pydict(_a			)			).func()
  
  
  
        class  lowerCAmelCase__   (       datasets.Dataset       ):
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          @fingerprint_transform(inplace=lowerCAmelCase_	, version="1.0.0"    )
          def 		_lowerCamelCase	(self  : Dict	, lowerCAmelCase_  : Any    )   ->       Any:
                return DummyDatasetChild(self.data	, fingerprint=lowerCAmelCase_    )
  
        snake_case_						       =							DummyDatasetChild(InMemoryTable.from_pydict(_a			)			).func()
  
  
  
        class  lowerCAmelCase__   (       datasets.Dataset       ):
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          @fingerprint_transform(inplace=lowerCAmelCase_	, version="2.0.0"    )
          def 		_lowerCamelCase	(self  : Union[str, Any]	, lowerCAmelCase_  : Any    )   ->       int:
                return DummyDatasetChild(self.data	, fingerprint=lowerCAmelCase_    )
  
        snake_case_						       =							DummyDatasetChild(InMemoryTable.from_pydict(_a			)			).func()
  
        assert len({fingeprint_no_version, fingeprint_a, fingeprint_a}			) == 3
  
  
  
  
  def   snake_case_       (							):
        # AttributeError: module 'dill._dill' has no attribute 'stack'
        snake_case_						       =							Hasher()
        hasher.update(lambda _a			: x			)
  
  
style_context: 
  '''simple docstring'''
  
  
  
  
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  def   snake_case_       (							_a     :							Optional[int] ,					_a     :							Dict ,					_a     :							Optional[Any] ,					_a     :							Optional[Any] ,					_a     :							int = 4_096			):
        snake_case_       ,	snake_case_       ,	snake_case_						       =							key.shape[-3:]
        snake_case_						       =							value.shape[-1]
        snake_case_						       =							min(_a ,					_a			)
        snake_case_						       =							query / jnp.sqrt(_a			)
  
        @functools.partial(jax.checkpoint ,					prevent_cse=_a			)
        def summarize_chunk(_a     :							Optional[int] ,					_a     :							List[str] ,					_a     :							Tuple			):
              snake_case_						       =							jnp.einsum("...qhd,...khd->...qhk" ,					_a ,					_a ,					precision=_a			)
  
              snake_case_						       =							jnp.max(_a ,					axis=-1 ,					keepdims=_a			)
              snake_case_						       =							jax.lax.stop_gradient(_a			)
              snake_case_						       =							jnp.exp(attn_weights - max_score			)
  
              snake_case_						       =							jnp.einsum("...vhf,...qhv->...qhf" ,					_a ,					_a ,					precision=_a			)
              snake_case_						       =							jnp.einsum("...qhk->...qh" ,					_a			)
  
              return (exp_values, exp_weights.sum(axis=-1			), max_score)
  
        def chunk_scanner(_a     :							Tuple			):
              # julienne key array
              snake_case_						       =							jax.lax.dynamic_slice(
                  operand=_a ,					start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0] ,					slice_sizes=list(key.shape[:-3]			) + [key_chunk_size, num_heads, k_features] ,					)
  
              # julienne value array
              snake_case_						       =							jax.lax.dynamic_slice(
                  operand=_a ,					start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0] ,					slice_sizes=list(value.shape[:-3]			) + [key_chunk_size, num_heads, v_features] ,					)
  
              return summarize_chunk(_a ,					_a ,					_a			)
  
        snake_case_       ,	snake_case_       ,	snake_case_						       =							jax.lax.map(f=_a ,					xs=jnp.arange(0 ,					_a ,					_a			)			)
  
        snake_case_						       =							jnp.max(_a ,					axis=0 ,					keepdims=_a			)
        snake_case_						       =							jnp.exp(chunk_max - global_max			)
  
        chunk_values *= jnp.expand_dims(_a ,					axis=-1			)
        chunk_weights *= max_diffs
  
        snake_case_						       =							chunk_values.sum(axis=0			)
        snake_case_						       =							jnp.expand_dims(_a ,					-1			).sum(axis=0			)
  
        return all_values / all_weights
  def   snake_case_       (							_a     :							Union[str, Any] ,					_a     :							Optional[int] ,					_a     :							List[Any] ,					_a     :							str=jax.lax.Precision.HIGHEST ,					_a     :							int = 1_024 ,					_a     :							int = 4_096			):
        snake_case_       ,	snake_case_       ,	snake_case_						       =							query.shape[-3:]
  
        def chunk_scanner(_a     :							Tuple ,					_a     :							Optional[int]			):
              # julienne query array
              snake_case_						       =							jax.lax.dynamic_slice(
                  operand=_a ,					start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0] ,					slice_sizes=list(query.shape[:-3]			) + [min(_a ,					_a			), num_heads, q_features] ,					)
  
              return (
                  chunk_idx + query_chunk_size,  # unused ignore it
                  _query_chunk_attention(
                      query=_a ,					key=_a ,					value=_a ,					precision=_a ,					key_chunk_size=_a			),
              )
  
        snake_case_       ,	snake_case_						       =							jax.lax.scan(
            f=_a ,					init=0 ,					xs=_a ,					length=math.ceil(num_q / query_chunk_size			)			)
  
        return jnp.concatenate(_a ,					axis=-3			)  # fuse the chunked result back
  
  
  
  class  lowerCAmelCase__   (       nn.Module       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    _UpperCamelCase						= 42
    _UpperCamelCase						= 8
    _UpperCamelCase						= 64
    _UpperCamelCase						= 0.0
    _UpperCamelCase						= False
    _UpperCamelCase						= jnp.floataa
  
    def 		_lowerCamelCase	(self  : Any    )   ->       Optional[Any]:
          snake_case_						       =							self.dim_head * self.heads
          snake_case_						       =							self.dim_head**-0.5
  
          # Weights were exported with old names {to_q, to_k, to_v, to_out}
          snake_case_						       =							nn.Dense(lowerCAmelCase_	, use_bias=lowerCAmelCase_	, dtype=self.dtype	, name="to_q"    )
          snake_case_						       =							nn.Dense(lowerCAmelCase_	, use_bias=lowerCAmelCase_	, dtype=self.dtype	, name="to_k"    )
          snake_case_						       =							nn.Dense(lowerCAmelCase_	, use_bias=lowerCAmelCase_	, dtype=self.dtype	, name="to_v"    )
  
          snake_case_						       =							nn.Dense(self.query_dim	, dtype=self.dtype	, name="to_out_0"    )
          snake_case_						       =							nn.Dropout(rate=self.dropout    )
  
    def 		_lowerCamelCase	(self  : Any	, lowerCAmelCase_  : Optional[Any]    )   ->       Tuple:
          snake_case_       ,	snake_case_       ,	snake_case_						       =							tensor.shape
          snake_case_						       =							self.heads
          snake_case_						       =							tensor.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, dim // head_size    )
          snake_case_						       =							jnp.transpose(lowerCAmelCase_	, (0, 2, 1, 3)    )
          snake_case_						       =							tensor.reshape(batch_size * head_size	, lowerCAmelCase_	, dim // head_size    )
          return tensor
  
    def 		_lowerCamelCase	(self  : List[Any]	, lowerCAmelCase_  : Optional[int]    )   ->       List[str]:
          snake_case_       ,	snake_case_       ,	snake_case_						       =							tensor.shape
          snake_case_						       =							self.heads
          snake_case_						       =							tensor.reshape(batch_size // head_size	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
          snake_case_						       =							jnp.transpose(lowerCAmelCase_	, (0, 2, 1, 3)    )
          snake_case_						       =							tensor.reshape(batch_size // head_size	, lowerCAmelCase_	, dim * head_size    )
          return tensor
  
  
    def __call__(self  : Tuple	, lowerCAmelCase_  : Optional[Any]	, lowerCAmelCase_  : Optional[Any]=None	, lowerCAmelCase_  : Union[str, Any]=True    )   ->       str:
          snake_case_						       =							hidden_states if context is None else context
  
          snake_case_						       =							self.query(lowerCAmelCase_    )
          snake_case_						       =							self.key(lowerCAmelCase_    )
          snake_case_						       =							self.value(lowerCAmelCase_    )
  
          snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
          snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
          snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
  
          if self.use_memory_efficient_attention:
                snake_case_						       =							query_states.transpose(1	, 0	, 2    )
                snake_case_						       =							key_states.transpose(1	, 0	, 2    )
                snake_case_						       =							value_states.transpose(1	, 0	, 2    )
  
                # this if statement create a chunk size for each layer of the unet
                # the chunk size is equal to the query_length dimension of the deepest layer of the unet
  
                snake_case_						       =							query_states.shape[-3]
                if flatten_latent_dim % 64 == 0:
                      snake_case_						       =							int(flatten_latent_dim / 64    )
                elif flatten_latent_dim % 16 == 0:
                      snake_case_						       =							int(flatten_latent_dim / 16    )
                elif flatten_latent_dim % 4 == 0:
                      snake_case_						       =							int(flatten_latent_dim / 4    )
                else:
                      snake_case_						       =							int(lowerCAmelCase_    )
  
                snake_case_						       =							jax_memory_efficient_attention(
                    lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, query_chunk_size=lowerCAmelCase_	, key_chunk_size=40_96 * 4    )
  
                snake_case_						       =							hidden_states.transpose(1	, 0	, 2    )
          else:
                # compute attentions
                snake_case_						       =							jnp.einsum("b i d, b j d->b i j"	, lowerCAmelCase_	, lowerCAmelCase_    )
                snake_case_						       =							attention_scores * self.scale
                snake_case_						       =							nn.softmax(lowerCAmelCase_	, axis=2    )
  
                # attend to values
                snake_case_						       =							jnp.einsum("b i j, b j d -> b i d"	, lowerCAmelCase_	, lowerCAmelCase_    )
  
          snake_case_						       =							self.reshape_batch_dim_to_heads(lowerCAmelCase_    )
          snake_case_						       =							self.proj_attn(lowerCAmelCase_    )
          return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
  
  
  
  class  lowerCAmelCase__   (       nn.Module       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    _UpperCamelCase						= 42
    _UpperCamelCase						= 42
    _UpperCamelCase						= 42
    _UpperCamelCase						= 0.0
    _UpperCamelCase						= False
    _UpperCamelCase						= jnp.floataa
    _UpperCamelCase						= False
  
    def 		_lowerCamelCase	(self  : str    )   ->       str:
          # self attention (or cross_attention if only_cross_attention is True)
          snake_case_						       =							FlaxAttention(
              self.dim	, self.n_heads	, self.d_head	, self.dropout	, self.use_memory_efficient_attention	, dtype=self.dtype    )
          # cross attention
          snake_case_						       =							FlaxAttention(
              self.dim	, self.n_heads	, self.d_head	, self.dropout	, self.use_memory_efficient_attention	, dtype=self.dtype    )
          snake_case_						       =							FlaxFeedForward(dim=self.dim	, dropout=self.dropout	, dtype=self.dtype    )
          snake_case_						       =							nn.LayerNorm(epsilon=1E-5	, dtype=self.dtype    )
          snake_case_						       =							nn.LayerNorm(epsilon=1E-5	, dtype=self.dtype    )
          snake_case_						       =							nn.LayerNorm(epsilon=1E-5	, dtype=self.dtype    )
          snake_case_						       =							nn.Dropout(rate=self.dropout    )
  
  
    def __call__(self  : Any	, lowerCAmelCase_  : Union[str, Any]	, lowerCAmelCase_  : Optional[Any]	, lowerCAmelCase_  : List[Any]=True    )   ->       Optional[int]:
          # self attention
          snake_case_						       =							hidden_states
          if self.only_cross_attention:
                snake_case_						       =							self.attna(self.norma(lowerCAmelCase_    )	, lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
          else:
                snake_case_						       =							self.attna(self.norma(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )
          snake_case_						       =							hidden_states + residual
  
          # cross attention
          snake_case_						       =							hidden_states
          snake_case_						       =							self.attna(self.norma(lowerCAmelCase_    )	, lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
          snake_case_						       =							hidden_states + residual
  
          # feed forward
          snake_case_						       =							hidden_states
          snake_case_						       =							self.ff(self.norma(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )
          snake_case_						       =							hidden_states + residual
  
          return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
  
  
  
  class  lowerCAmelCase__   (       nn.Module       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    _UpperCamelCase						= 42
    _UpperCamelCase						= 42
    _UpperCamelCase						= 42
    _UpperCamelCase						= 1
    _UpperCamelCase						= 0.0
    _UpperCamelCase						= False
    _UpperCamelCase						= False
    _UpperCamelCase						= jnp.floataa
    _UpperCamelCase						= False
  
    def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Tuple:
          snake_case_						       =							nn.GroupNorm(num_groups=32	, epsilon=1E-5    )
  
          snake_case_						       =							self.n_heads * self.d_head
          if self.use_linear_projection:
                snake_case_						       =							nn.Dense(lowerCAmelCase_	, dtype=self.dtype    )
          else:
                snake_case_						       =							nn.Conv(
                    lowerCAmelCase_	, kernel_size=(1, 1)	, strides=(1, 1)	, padding="VALID"	, dtype=self.dtype	, )
  
          snake_case_						       =							[
              FlaxBasicTransformerBlock(
                  lowerCAmelCase_	, self.n_heads	, self.d_head	, dropout=self.dropout	, only_cross_attention=self.only_cross_attention	, dtype=self.dtype	, use_memory_efficient_attention=self.use_memory_efficient_attention	, )
              for _ in range(self.depth    )
          ]
  
          if self.use_linear_projection:
                snake_case_						       =							nn.Dense(lowerCAmelCase_	, dtype=self.dtype    )
          else:
                snake_case_						       =							nn.Conv(
                    lowerCAmelCase_	, kernel_size=(1, 1)	, strides=(1, 1)	, padding="VALID"	, dtype=self.dtype	, )
  
          snake_case_						       =							nn.Dropout(rate=self.dropout    )
  
  
    def __call__(self  : int	, lowerCAmelCase_  : List[Any]	, lowerCAmelCase_  : str	, lowerCAmelCase_  : str=True    )   ->       Optional[Any]:
          snake_case_       ,	snake_case_       ,	snake_case_       ,	snake_case_						       =							hidden_states.shape
          snake_case_						       =							hidden_states
          snake_case_						       =							self.norm(lowerCAmelCase_    )
          if self.use_linear_projection:
                snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, height * width	, lowerCAmelCase_    )
                snake_case_						       =							self.proj_in(lowerCAmelCase_    )
          else:
                snake_case_						       =							self.proj_in(lowerCAmelCase_    )
                snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, height * width	, lowerCAmelCase_    )
  
          for transformer_block in self.transformer_blocks:
                snake_case_						       =							transformer_block(lowerCAmelCase_	, lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
  
          if self.use_linear_projection:
                snake_case_						       =							self.proj_out(lowerCAmelCase_    )
                snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
          else:
                snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
                snake_case_						       =							self.proj_out(lowerCAmelCase_    )
  
          snake_case_						       =							hidden_states + residual
          return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
  
  
  
  class  lowerCAmelCase__   (       nn.Module       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    _UpperCamelCase						= 42
    _UpperCamelCase						= 0.0
    _UpperCamelCase						= jnp.floataa
  
    def 		_lowerCamelCase	(self  : str    )   ->       List[str]:
          # The second linear layer needs to be called
          # net_2 for now to match the index of the Sequential layer
          snake_case_						       =							FlaxGEGLU(self.dim	, self.dropout	, self.dtype    )
          snake_case_						       =							nn.Dense(self.dim	, dtype=self.dtype    )
  
  
    def __call__(self  : int	, lowerCAmelCase_  : Tuple	, lowerCAmelCase_  : Any=True    )   ->       int:
          snake_case_						       =							self.net_a(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
          snake_case_						       =							self.net_a(lowerCAmelCase_    )
          return hidden_states
  
  
  
  class  lowerCAmelCase__   (       nn.Module       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    _UpperCamelCase						= 42
    _UpperCamelCase						= 0.0
    _UpperCamelCase						= jnp.floataa
  
    def 		_lowerCamelCase	(self  : str    )   ->       Optional[Any]:
          snake_case_						       =							self.dim * 4
          snake_case_						       =							nn.Dense(inner_dim * 2	, dtype=self.dtype    )
          snake_case_						       =							nn.Dropout(rate=self.dropout    )
  
  
    def __call__(self  : Tuple	, lowerCAmelCase_  : Dict	, lowerCAmelCase_  : Optional[Any]=True    )   ->       Union[str, Any]:
          snake_case_						       =							self.proj(lowerCAmelCase_    )
          snake_case_       ,	snake_case_						       =							jnp.split(lowerCAmelCase_	, 2	, axis=2    )
          return self.dropout_layer(hidden_linear * nn.gelu(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )
  
  