






"""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 						_a  :





   '''simple docstring'''






   def __init__(       self					, UpperCAmelCase_)			->   List[Any]:




          '''simple docstring'''




          lowercase__:							List[str]         =     foo





   def __call__(       self)			->   Union[str, Any]:




          '''simple docstring'''




          return self.foo







class 						_a  (		datasets.Dataset       ):





   '''simple docstring'''






   @fingerprint_transform(inplace=UpperCAmelCase_)
   def __lowercase							(       self					, UpperCAmelCase_					, *UpperCAmelCase_					, **UpperCAmelCase_)			->   int:




          '''simple docstring'''




          return DatasetChild(self.data					, fingerprint=UpperCAmelCase_)





   @fingerprint_transform(inplace=UpperCAmelCase_)
   def __lowercase							(       self					, UpperCAmelCase_					, *UpperCAmelCase_					, **UpperCAmelCase_)			->   int:




          '''simple docstring'''




          return DatasetChild(self.data					, fingerprint=UpperCAmelCase_)







class 						_a  :





   '''simple docstring'''






   def __init__(       self					, UpperCAmelCase_)			->   Optional[int]:




          '''simple docstring'''




          lowercase__:							Tuple         =     callable





   def __call__(       self					, *UpperCAmelCase_					, **UpperCAmelCase_)			->   Optional[int]:




          '''simple docstring'''




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





   def __getstate__(       self)			->   int:




          '''simple docstring'''




          raise pickle.PicklingError()







class 						_a  (		lowercase_       ):





   '''simple docstring'''






   @require_transformers
   @pytest.mark.integration
   def __lowercase							(       self)			->   List[Any]:




          '''simple docstring'''




          from transformers import AutoTokenizer

          def encode(UpperCAmelCase_):
                 return tokenizer(UpperCAmelCase_)

          # TODO: add hash consistency tests across sessions
          lowercase__:							Tuple         =     AutoTokenizer.from_pretrained("bert-base-uncased")
          lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(lambda UpperCAmelCase_: tokenizer(UpperCAmelCase_))).hexdigest()
          lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Tuple         =     AutoTokenizer.from_pretrained("bert-base-cased")
          lowercase__:							int         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Tuple         =     mda(datasets.utils.py_utils.dumps(lambda UpperCAmelCase_: tokenizer(UpperCAmelCase_))).hexdigest()
          lowercase__:							List[str]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Optional[int]         =     AutoTokenizer.from_pretrained("bert-base-uncased")
          lowercase__:							Tuple         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							str         =     mda(datasets.utils.py_utils.dumps(lambda UpperCAmelCase_: tokenizer(UpperCAmelCase_))).hexdigest()
          lowercase__:							Tuple         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   @require_transformers
   @pytest.mark.integration
   def __lowercase							(       self)			->   Optional[int]:




          '''simple docstring'''




          from transformers import AutoTokenizer

          lowercase__:							Optional[int]         =     AutoTokenizer.from_pretrained("gpt2")
          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          tokenizer("Hello world !")  # call once to change the tokenizer's cache
          lowercase__:							Tuple         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)





   @require_regex
   def __lowercase							(       self)			->   str:




          '''simple docstring'''




          import regex

          lowercase__:							List[Any]         =     regex.Regex("foo")
          lowercase__:							Optional[int]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							List[Any]         =     regex.Regex("bar")
          lowercase__:							List[str]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Union[str, Any]         =     regex.Regex("foo")
          lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)







class 						_a  (		lowercase_       ):





   '''simple docstring'''






   def __lowercase							(       self)			->   int:




          '''simple docstring'''




          def func():
                 return foo

          lowercase__:							int         =     [0]
          lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Dict         =     [1]
          lowercase__:							int         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Union[str, Any]         =     [0]
          lowercase__:							List[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Dict:




          '''simple docstring'''




          def func():
                 pass

          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()

          def func():
                 pass

          lowercase__:							List[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   str:




          '''simple docstring'''




          lowercase__:							Optional[int]         =     mda(datasets.utils.py_utils.dumps(Foo([0]))).hexdigest()
          lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(Foo([1]))).hexdigest()
          lowercase__:							int         =     mda(datasets.utils.py_utils.dumps(Foo([0]))).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Union[str, Any]:




          '''simple docstring'''




          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(Foo([0]).__call__)).hexdigest()
          lowercase__:							List[Any]         =     mda(datasets.utils.py_utils.dumps(Foo([1]).__call__)).hexdigest()
          lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(Foo([0]).__call__)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Union[str, Any]:




          '''simple docstring'''




          lowercase__:							List[Any]         =     (
              "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(*UpperCAmelCase_):
                 if hasattr(args[-3]					, "encode"):
                        lowercase__:							Tuple         =     list(UpperCAmelCase_)
                        lowercase__:							str         =     args[-3].encode()  # co_lnotab
                        lowercase__:							Any         =     args[-10].encode()  # co_code
                 if hasattr(UpperCAmelCase_					, "co_posonlyargcount"):
                        if len(UpperCAmelCase_) == 16:
                               return CodeType(*UpperCAmelCase_)
                        elif len(UpperCAmelCase_) == 15:
                               return CodeType(args[0]					, 0					, *args[1:])
                        return CodeType(args[0]					, 0					, 0					, *args[1:])
                 elif hasattr(UpperCAmelCase_					, "co_kwonlyargcount"):
                        if len(UpperCAmelCase_) == 16:
                               return CodeType(args[0]					, *args[2:])
                        elif len(UpperCAmelCase_) == 15:
                               return CodeType(*UpperCAmelCase_)
                        return CodeType(args[0]					, 0					, *args[1:])
                 if len(UpperCAmelCase_) == 16:
                        return CodeType(args[0]					, *args[3:])
                 elif len(UpperCAmelCase_) == 15:
                        return CodeType(args[0]					, *args[2:])
                 return CodeType(*UpperCAmelCase_)

          def create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_):
                 def func():
                        return returned_obj

                 lowercase__:							Dict         =     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"):
                        lowercase__:							Union[str, Any]         =     _create_code_pyaa(
                            *[getattr(UpperCAmelCase_					, UpperCAmelCase_) if k != "co_filename" else co_filename for k in code_args_pyaa])
                 else:
                        lowercase__:							Tuple         =     code.replace(co_filename=UpperCAmelCase_)
                 return FunctionType(UpperCAmelCase_					, func.__globals__					, func.__name__					, func.__defaults__					, func.__closure__)

          lowercase__    ,      lowercase__:							Any         =     "<ipython-input-2-e0383a102aae>", [0]
          lowercase__:							List[str]         =     mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_))).hexdigest()
          lowercase__    ,      lowercase__:							Optional[Any]         =     "<ipython-input-2-e0383a102aae>", [1]
          lowercase__:							Optional[int]         =     mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_))).hexdigest()
          lowercase__    ,      lowercase__:							Tuple         =     "<ipython-input-5-713f6613acf3>", [0]
          lowercase__:							List[str]         =     mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_))).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)

          lowercase__    ,      lowercase__:							Dict         =     os.path.join(gettempdir()					, "ipykernel_12345"					, "321456789.py"), [0]
          lowercase__:							Tuple         =     mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_))).hexdigest()
          lowercase__    ,      lowercase__:							Dict         =     os.path.join(gettempdir()					, "ipykernel_12345"					, "321456789.py"), [1]
          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_))).hexdigest()
          lowercase__    ,      lowercase__:							Any         =     os.path.join(gettempdir()					, "ipykernel_12345"					, "654123987.py"), [0]
          lowercase__:							Dict         =     mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase_					, UpperCAmelCase_))).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Dict:




          '''simple docstring'''




          lowercase__    ,      lowercase__:							str         =     [0], [1]

          def func():
                 return foo, bar

          lowercase__:							List[str]         =     "__main__"

          def globalvars_mocka_side_effect(UpperCAmelCase_					, *UpperCAmelCase_					, **UpperCAmelCase_):
                 return {"foo": foo, "bar": bar}

          def globalvars_mocka_side_effect(UpperCAmelCase_					, *UpperCAmelCase_					, **UpperCAmelCase_):
                 return {"bar": bar, "foo": foo}

          with patch("dill.detect.globalvars"					, side_effect=UpperCAmelCase_) as globalvars_mocka:
                 lowercase__:							Optional[Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
                 self.assertGreater(globalvars_mocka.call_count					, 0)
          with patch("dill.detect.globalvars"					, side_effect=UpperCAmelCase_) as globalvars_mocka:
                 lowercase__:							Tuple         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
                 self.assertGreater(globalvars_mocka.call_count					, 0)
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)







class 						_a  (		lowercase_       ):





   '''simple docstring'''






   def __lowercase							(       self)			->   int:




          '''simple docstring'''




          lowercase__:							List[Any]         =     Hasher.hash("hello")
          lowercase__:							Optional[int]         =     Hasher.hash("hello")
          lowercase__:							int         =     Hasher.hash("there")
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Any:




          '''simple docstring'''




          lowercase__:							str         =     Hasher.hash(Foo("hello"))
          lowercase__:							Tuple         =     Hasher.hash(Foo("hello"))
          lowercase__:							Union[str, Any]         =     Hasher.hash(Foo("there"))
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   str:




          '''simple docstring'''




          lowercase__:							str         =     Hasher()
          for x in ["hello", Foo("hello")]:
                 hasher.update(UpperCAmelCase_)
          lowercase__:							Dict         =     hasher.hexdigest()
          lowercase__:							Any         =     Hasher()
          for x in ["hello", Foo("hello")]:
                 hasher.update(UpperCAmelCase_)
          lowercase__:							List[Any]         =     hasher.hexdigest()
          lowercase__:							Optional[int]         =     Hasher()
          for x in ["there", Foo("there")]:
                 hasher.update(UpperCAmelCase_)
          lowercase__:							Tuple         =     hasher.hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   List[str]:




          '''simple docstring'''




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





   def __lowercase							(       self)			->   Optional[Any]:




          '''simple docstring'''




          lowercase__:							Optional[Any]         =     "abc"
          lowercase__:							Any         =     [string, string]  # two strings have the same ids
          lowercase__:							str         =     [string, string]
          lowercase__:							Optional[int]         =     json.loads(F"""[\"{string}\", \"{string}\"]""")  # two strings have different ids
          self.assertIs(obja[0]					, UpperCAmelCase_)
          self.assertIs(obja[0]					, obja[1])
          self.assertIs(obja[0]					, UpperCAmelCase_)
          self.assertIs(obja[0]					, obja[1])
          self.assertIsNot(obja[0]					, UpperCAmelCase_)
          self.assertIsNot(obja[0]					, obja[1])
          lowercase__:							List[Any]         =     Hasher.hash(UpperCAmelCase_)
          lowercase__:							str         =     Hasher.hash(UpperCAmelCase_)
          lowercase__:							Tuple         =     Hasher.hash(UpperCAmelCase_)
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)





   @require_tiktoken
   def __lowercase							(       self)			->   str:




          '''simple docstring'''




          import tiktoken

          lowercase__:							Union[str, Any]         =     tiktoken.get_encoding("gpt2")
          lowercase__:							List[str]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Dict         =     tiktoken.get_encoding("r50k_base")
          lowercase__:							Dict         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							List[Any]         =     tiktoken.get_encoding("gpt2")
          lowercase__:							Optional[int]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   @require_torch
   def __lowercase							(       self)			->   List[Any]:




          '''simple docstring'''




          import torch

          lowercase__:							Tuple         =     torch.tensor([1.0])
          lowercase__:							Optional[int]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Dict         =     torch.tensor([2.0])
          lowercase__:							Union[str, Any]         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Tuple         =     torch.tensor([1.0])
          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)





   @require_spacy
   @require_spacy_model("en_core_web_sm")
   @require_spacy_model("fr_core_news_sm")
   @pytest.mark.integration
   def __lowercase							(       self)			->   Any:




          '''simple docstring'''




          import spacy

          lowercase__:							Tuple         =     spacy.load("en_core_web_sm")
          lowercase__:							Any         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							Any         =     spacy.load("fr_core_news_sm")
          lowercase__:							str         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          lowercase__:							int         =     spacy.load("en_core_web_sm")
          lowercase__:							str         =     mda(datasets.utils.py_utils.dumps(UpperCAmelCase_)).hexdigest()
          self.assertEqual(UpperCAmelCase_					, UpperCAmelCase_)
          self.assertNotEqual(UpperCAmelCase_					, UpperCAmelCase_)







@pytest.mark.integration
def 							A(    snake_case_						):
       """simple docstring"""





       lowercase__:							Optional[Any]         =     tmp_path / "dir1"
       lowercase__:							Tuple         =     tmp_path / "dir2"
       dira.mkdir()
       dira.mkdir()
       lowercase__:							Union[str, Any]         =     "script.py"
       lowercase__:							Any         =     dedent(
           "\n    from datasets.fingerprint import Hasher\n    def foo():\n        pass\n    print(Hasher.hash(foo))\n    "						)
       lowercase__:							List[str]         =     dira / script_filename
       lowercase__:							int         =     dira / script_filename
       with script_patha.open("w"						) as f:
              f.write(snake_case_						)
       with script_patha.open("w"						) as f:
              f.write(snake_case_						)
       lowercase__:							Any         =     subprocess.check_output(["python", str(snake_case_						)]						)
       lowercase__:							Optional[Any]         =     subprocess.check_output(["python", str(snake_case_						)]						)
       assert fingerprinta == fingerprinta






def 							A(    ):
       """simple docstring"""





       lowercase__:							Optional[int]         =     {"a": [0, 1, 2]}
       lowercase__:							Dict         =     DatasetChild(InMemoryTable.from_pydict(snake_case_						)						)
       lowercase__:							List[Any]         =     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 							A(    ):
       """simple docstring"""





       lowercase__:							List[str]         =     {"a": [0, 1, 2]}







       class 						_a  (		datasets.Dataset       ):





          '''simple docstring'''






          @fingerprint_transform(inplace=UpperCAmelCase_)
          def __lowercase							(       self					, UpperCAmelCase_)			->   List[str]:




                 '''simple docstring'''




                 return DummyDatasetChild(self.data					, fingerprint=UpperCAmelCase_)







       lowercase__:							Tuple         =     DummyDatasetChild(InMemoryTable.from_pydict(snake_case_						)						).func()







       class 						_a  (		datasets.Dataset       ):





          '''simple docstring'''






          @fingerprint_transform(inplace=UpperCAmelCase_					, version="1.0.0")
          def __lowercase							(       self					, UpperCAmelCase_)			->   Optional[int]:




                 '''simple docstring'''




                 return DummyDatasetChild(self.data					, fingerprint=UpperCAmelCase_)







       lowercase__:							Optional[int]         =     DummyDatasetChild(InMemoryTable.from_pydict(snake_case_						)						).func()







       class 						_a  (		datasets.Dataset       ):





          '''simple docstring'''






          @fingerprint_transform(inplace=UpperCAmelCase_					, version="2.0.0")
          def __lowercase							(       self					, UpperCAmelCase_)			->   Any:




                 '''simple docstring'''




                 return DummyDatasetChild(self.data					, fingerprint=UpperCAmelCase_)







       lowercase__:							Optional[int]         =     DummyDatasetChild(InMemoryTable.from_pydict(snake_case_						)						).func()

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







def 							A(    ):
       """simple docstring"""





       lowercase__:							int         =     Hasher()
       hasher.update(lambda snake_case_						: x						)













"""simple docstring"""




import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase = 4096   ):






  """simple docstring"""





  A_		,			A_		,			A_    :			List[Any]									=    key.shape[-3:]
  A_    :			Dict									=    value.shape[-1]
  A_    :			Union[str, Any]									=    min(_UpperCAmelCase					,   _UpperCAmelCase   )
  A_    :			str									=    query / jnp.sqrt(_UpperCAmelCase   )

  @functools.partial(jax.checkpoint					,   prevent_cse=_UpperCAmelCase   )
  def summarize_chunk(_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase   ):
    A_    :			Any									=    jnp.einsum('...qhd,...khd->...qhk'					,   _UpperCAmelCase					,   _UpperCAmelCase					,   precision=_UpperCAmelCase   )

    A_    :			int									=    jnp.max(_UpperCAmelCase					,   axis=-1					,   keepdims=_UpperCAmelCase   )
    A_    :			List[str]									=    jax.lax.stop_gradient(_UpperCAmelCase   )
    A_    :			Dict									=    jnp.exp(attn_weights - max_score   )

    A_    :			int									=    jnp.einsum('...vhf,...qhv->...qhf'					,   _UpperCAmelCase					,   _UpperCAmelCase					,   precision=_UpperCAmelCase   )
    A_    :			Optional[Any]									=    jnp.einsum('...qhk->...qh'					,   _UpperCAmelCase   )

    return (exp_values, exp_weights.sum(axis=-1   ), max_score)

  def chunk_scanner(_UpperCAmelCase   ):
    # julienne key array
    A_    :			Dict									=    jax.lax.dynamic_slice(
        operand=_UpperCAmelCase					,   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
    A_    :			Optional[Any]									=    jax.lax.dynamic_slice(
        operand=_UpperCAmelCase					,   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(_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase   )

  A_		,			A_		,			A_    :			Dict									=    jax.lax.map(f=_UpperCAmelCase					,   xs=jnp.arange(0					,   _UpperCAmelCase					,   _UpperCAmelCase   )   )

  A_    :			int									=    jnp.max(_UpperCAmelCase					,   axis=0					,   keepdims=_UpperCAmelCase   )
  A_    :			Tuple									=    jnp.exp(chunk_max - global_max   )

  chunk_values *= jnp.expand_dims(_UpperCAmelCase					,   axis=-1   )
  chunk_weights *= max_diffs

  A_    :			List[Any]									=    chunk_values.sum(axis=0   )
  A_    :			Dict									=    jnp.expand_dims(_UpperCAmelCase					,   -1   ).sum(axis=0   )

  return all_values / all_weights





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase=jax.lax.Precision.HIGHEST					,   _UpperCAmelCase = 1024					,   _UpperCAmelCase = 4096   ):






  """simple docstring"""





  A_		,			A_		,			A_    :			Any									=    query.shape[-3:]

  def chunk_scanner(_UpperCAmelCase					,   _UpperCAmelCase   ):
    # julienne query array
    A_    :			Tuple									=    jax.lax.dynamic_slice(
        operand=_UpperCAmelCase					,   start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]					,   slice_sizes=list(query.shape[:-3]   ) + [min(_UpperCAmelCase					,   _UpperCAmelCase   ), num_heads, q_features]					,   )

    return (
        chunk_idx + query_chunk_size,  # unused ignore it
        _query_chunk_attention(
            query=_UpperCAmelCase					,   key=_UpperCAmelCase					,   value=_UpperCAmelCase					,   precision=_UpperCAmelCase					,   key_chunk_size=_UpperCAmelCase   ),
    )

  A_		,			A_    :			Union[str, Any]									=    jax.lax.scan(
      f=_UpperCAmelCase					,   init=0					,   xs=_UpperCAmelCase					,   length=math.ceil(num_q / query_chunk_size   )   )

  return jnp.concatenate(_UpperCAmelCase					,   axis=-3   )  # fuse the chunked result back

class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''
      lowercase_ :					int
      lowercase_ :					int             =  8
      lowercase_ :					int             =  64
      lowercase_ :					float             =  0.0
      lowercase_ :					bool             =  False
      lowercase_ :					jnp.dtype             =  jnp.floataa

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[str]									=    self.dim_head * self.heads
        A_    :			Tuple									=    self.dim_head**-0.5

        # Weights were exported with old names {to_q, to_k, to_v, to_out}
        A_    :			Any									=    nn.Dense(snake_case_   , use_bias=snake_case_   , dtype=self.dtype   , name='to_q'    )
        A_    :			int									=    nn.Dense(snake_case_   , use_bias=snake_case_   , dtype=self.dtype   , name='to_k'    )
        A_    :			Optional[int]									=    nn.Dense(snake_case_   , use_bias=snake_case_   , dtype=self.dtype   , name='to_v'    )

        A_    :			Union[str, Any]									=    nn.Dense(self.query_dim   , dtype=self.dtype   , name='to_out_0'    )
        A_    :			str									=    nn.Dropout(rate=self.dropout    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_		,			A_		,			A_    :			Union[str, Any]									=    tensor.shape
        A_    :			List[str]									=    self.heads
        A_    :			str									=    tensor.reshape(snake_case_   , snake_case_   , snake_case_   , dim // head_size    )
        A_    :			Optional[Any]									=    jnp.transpose(snake_case_   , (0, 2, 1, 3)    )
        A_    :			Any									=    tensor.reshape(batch_size * head_size   , snake_case_   , dim // head_size    )
        return tensor

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_		,			A_		,			A_    :			str									=    tensor.shape
        A_    :			Optional[Any]									=    self.heads
        A_    :			List[Any]									=    tensor.reshape(batch_size // head_size   , snake_case_   , snake_case_   , snake_case_    )
        A_    :			List[Any]									=    jnp.transpose(snake_case_   , (0, 2, 1, 3)    )
        A_    :			Any									=    tensor.reshape(batch_size // head_size   , snake_case_   , dim * head_size    )
        return tensor

      def __call__(	self   , snake_case_   , snake_case_=None   , snake_case_=True    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    hidden_states if context is None else context

        A_    :			str									=    self.query(snake_case_    )
        A_    :			Any									=    self.key(snake_case_    )
        A_    :			Union[str, Any]									=    self.value(snake_case_    )

        A_    :			int									=    self.reshape_heads_to_batch_dim(snake_case_    )
        A_    :			List[str]									=    self.reshape_heads_to_batch_dim(snake_case_    )
        A_    :			str									=    self.reshape_heads_to_batch_dim(snake_case_    )

        if self.use_memory_efficient_attention:
          A_    :			str									=    query_states.transpose(1   , 0   , 2    )
          A_    :			Optional[int]									=    key_states.transpose(1   , 0   , 2    )
          A_    :			Optional[Any]									=    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

          A_    :			Union[str, Any]									=    query_states.shape[-3]
          if flatten_latent_dim % 6_4 == 0:
            A_    :			Any									=    int(flatten_latent_dim / 6_4    )
          elif flatten_latent_dim % 1_6 == 0:
            A_    :			Optional[int]									=    int(flatten_latent_dim / 1_6    )
          elif flatten_latent_dim % 4 == 0:
            A_    :			Dict									=    int(flatten_latent_dim / 4    )
          else:
            A_    :			Optional[int]									=    int(snake_case_    )

          A_    :			Any									=    jax_memory_efficient_attention(
              snake_case_   , snake_case_   , snake_case_   , query_chunk_size=snake_case_   , key_chunk_size=4_0_9_6 * 4    )

          A_    :			Any									=    hidden_states.transpose(1   , 0   , 2    )
        else:
          # compute attentions
          A_    :			str									=    jnp.einsum('b i d, b j d->b i j'   , snake_case_   , snake_case_    )
          A_    :			Dict									=    attention_scores * self.scale
          A_    :			List[str]									=    nn.softmax(snake_case_   , axis=2    )

          # attend to values
          A_    :			Optional[Any]									=    jnp.einsum('b i j, b j d -> b i d'   , snake_case_   , snake_case_    )

        A_    :			List[str]									=    self.reshape_batch_dim_to_heads(snake_case_    )
        A_    :			Tuple									=    self.proj_attn(snake_case_    )
        return self.dropout_layer(snake_case_   , deterministic=snake_case_    )


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''
      lowercase_ :					int
      lowercase_ :					int
      lowercase_ :					int
      lowercase_ :					float             =  0.0
      lowercase_ :					bool             =  False
      lowercase_ :					jnp.dtype             =  jnp.floataa
      lowercase_ :					bool             =  False

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[str]									=    FlaxAttention(
            self.dim   , self.n_heads   , self.d_head   , self.dropout   , self.use_memory_efficient_attention   , dtype=self.dtype    )
        # cross attention
        A_    :			Optional[Any]									=    FlaxAttention(
            self.dim   , self.n_heads   , self.d_head   , self.dropout   , self.use_memory_efficient_attention   , dtype=self.dtype    )
        A_    :			Optional[Any]									=    FlaxFeedForward(dim=self.dim   , dropout=self.dropout   , dtype=self.dtype    )
        A_    :			Any									=    nn.LayerNorm(epsilon=1E-5   , dtype=self.dtype    )
        A_    :			Dict									=    nn.LayerNorm(epsilon=1E-5   , dtype=self.dtype    )
        A_    :			Optional[int]									=    nn.LayerNorm(epsilon=1E-5   , dtype=self.dtype    )
        A_    :			List[str]									=    nn.Dropout(rate=self.dropout    )

      def __call__(	self   , snake_case_   , snake_case_   , snake_case_=True    ):

        """simple docstring"""

        A_    :			Optional[int]									=    hidden_states
        if self.only_cross_attention:
          A_    :			Union[str, Any]									=    self.attna(self.norma(snake_case_    )   , snake_case_   , deterministic=snake_case_    )
        else:
          A_    :			Tuple									=    self.attna(self.norma(snake_case_    )   , deterministic=snake_case_    )
        A_    :			Optional[Any]									=    hidden_states + residual

        # cross attention
        A_    :			List[str]									=    hidden_states
        A_    :			List[str]									=    self.attna(self.norma(snake_case_    )   , snake_case_   , deterministic=snake_case_    )
        A_    :			Any									=    hidden_states + residual

        # feed forward
        A_    :			Dict									=    hidden_states
        A_    :			str									=    self.ff(self.norma(snake_case_    )   , deterministic=snake_case_    )
        A_    :			Optional[int]									=    hidden_states + residual

        return self.dropout_layer(snake_case_   , deterministic=snake_case_    )


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''
      lowercase_ :					int
      lowercase_ :					int
      lowercase_ :					int
      lowercase_ :					int             =  1
      lowercase_ :					float             =  0.0
      lowercase_ :					bool             =  False
      lowercase_ :					bool             =  False
      lowercase_ :					jnp.dtype             =  jnp.floataa
      lowercase_ :					bool             =  False

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    nn.GroupNorm(num_groups=3_2   , epsilon=1E-5    )

        A_    :			Dict									=    self.n_heads * self.d_head
        if self.use_linear_projection:
          A_    :			Union[str, Any]									=    nn.Dense(snake_case_   , dtype=self.dtype    )
        else:
          A_    :			Any									=    nn.Conv(
              snake_case_   , kernel_size=(1, 1)   , strides=(1, 1)   , padding='VALID'   , dtype=self.dtype   , )

        A_    :			List[str]									=    [
            FlaxBasicTransformerBlock(
                snake_case_   , 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:
          A_    :			List[Any]									=    nn.Dense(snake_case_   , dtype=self.dtype    )
        else:
          A_    :			Dict									=    nn.Conv(
              snake_case_   , kernel_size=(1, 1)   , strides=(1, 1)   , padding='VALID'   , dtype=self.dtype   , )

        A_    :			List[str]									=    nn.Dropout(rate=self.dropout    )

      def __call__(	self   , snake_case_   , snake_case_   , snake_case_=True    ):

        """simple docstring"""

        A_		,			A_		,			A_		,			A_    :			Any									=    hidden_states.shape
        A_    :			Any									=    hidden_states
        A_    :			List[Any]									=    self.norm(snake_case_    )
        if self.use_linear_projection:
          A_    :			str									=    hidden_states.reshape(snake_case_   , height * width   , snake_case_    )
          A_    :			Optional[Any]									=    self.proj_in(snake_case_    )
        else:
          A_    :			List[str]									=    self.proj_in(snake_case_    )
          A_    :			List[Any]									=    hidden_states.reshape(snake_case_   , height * width   , snake_case_    )

        for transformer_block in self.transformer_blocks:
          A_    :			int									=    transformer_block(snake_case_   , snake_case_   , deterministic=snake_case_    )

        if self.use_linear_projection:
          A_    :			Union[str, Any]									=    self.proj_out(snake_case_    )
          A_    :			Any									=    hidden_states.reshape(snake_case_   , snake_case_   , snake_case_   , snake_case_    )
        else:
          A_    :			int									=    hidden_states.reshape(snake_case_   , snake_case_   , snake_case_   , snake_case_    )
          A_    :			Union[str, Any]									=    self.proj_out(snake_case_    )

        A_    :			Union[str, Any]									=    hidden_states + residual
        return self.dropout_layer(snake_case_   , deterministic=snake_case_    )


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''
      lowercase_ :					int
      lowercase_ :					float             =  0.0
      lowercase_ :					jnp.dtype             =  jnp.floataa

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[str]									=    FlaxGEGLU(self.dim   , self.dropout   , self.dtype    )
        A_    :			int									=    nn.Dense(self.dim   , dtype=self.dtype    )

      def __call__(	self   , snake_case_   , snake_case_=True    ):

        """simple docstring"""

        A_    :			Dict									=    self.net_a(snake_case_   , deterministic=snake_case_    )
        A_    :			Any									=    self.net_a(snake_case_    )
        return hidden_states


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''
      lowercase_ :					int
      lowercase_ :					float             =  0.0
      lowercase_ :					jnp.dtype             =  jnp.floataa

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[int]									=    self.dim * 4
        A_    :			List[str]									=    nn.Dense(inner_dim * 2   , dtype=self.dtype    )
        A_    :			int									=    nn.Dropout(rate=self.dropout    )

      def __call__(	self   , snake_case_   , snake_case_=True    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    self.proj(snake_case_    )
        A_		,			A_    :			Any									=    jnp.split(snake_case_   , 2   , axis=2    )
        return self.dropout_layer(hidden_linear * nn.gelu(snake_case_    )   , deterministic=snake_case_    )