




import os
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from huggingface_hub.file_download import http_get
from requests.exceptions import HTTPError

from transformers import (
    AlbertTokenizer,
    AutoTokenizer,
    BertTokenizer,
    BertTokenizerFast,
    GPTaTokenizerFast,
    is_tokenizers_available,
)
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_tokenizers
from transformers.tokenization_utils import Trie


sys.path.append(str(Path(__file__).parent.parent / '''utils'''))

from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
   from test_module.custom_tokenization_fast import CustomTokenizerFast



class        a					(			unittest.TestCase   ):





     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        __lowerCamelCase:				Dict     =	mock.Mock()
        __lowerCamelCase:				List[str]     =	500
        __lowerCamelCase:				List[str]     =	{}
        __lowerCamelCase:				Union[str, Any]     =	HTTPError
        __lowerCamelCase:				Any     =	{}

        # Download this model to make sure it's in the cache.
        __lowerCamelCase:				int     =	BertTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert"""							)

        # Under the mock environment we get a 500 error when trying to reach the tokenizer.
        with mock.patch("""requests.Session.request"""    ,	return_value=SCREAMING_SNAKE_CASE_							) as mock_head:
           __lowerCamelCase:				Any     =	BertTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert"""							)
           # This check we did call the fake head request
           mock_head.assert_called()





     @require_tokenizers
     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        __lowerCamelCase:				Any     =	mock.Mock()
        __lowerCamelCase:				Dict     =	500
        __lowerCamelCase:				str     =	{}
        __lowerCamelCase:				Tuple     =	HTTPError
        __lowerCamelCase:				Optional[int]     =	{}

        # Download this model to make sure it's in the cache.
        __lowerCamelCase:				int     =	GPTaTokenizerFast.from_pretrained("""gpt2"""							)

        # Under the mock environment we get a 500 error when trying to reach the tokenizer.
        with mock.patch("""requests.Session.request"""    ,	return_value=SCREAMING_SNAKE_CASE_							) as mock_head:
           __lowerCamelCase:				Union[str, Any]     =	GPTaTokenizerFast.from_pretrained("""gpt2"""							)
           # This check we did call the fake head request
           mock_head.assert_called()





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):







        try:
           __lowerCamelCase:				Dict     =	tempfile.mktemp()
           with open(SCREAMING_SNAKE_CASE_    ,	"""wb"""							) as f:
              http_get("""https://huggingface.co/albert-base-v1/resolve/main/spiece.model"""    ,	SCREAMING_SNAKE_CASE_							)

           __lowerCamelCase:				Any     =	AlbertTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_							)
        finally:
           os.remove(SCREAMING_SNAKE_CASE_							)

        # Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
        # the current folder and have the right name.
        if os.path.isfile("""tokenizer.json"""							):
           # We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
           return
        try:
           with open("""tokenizer.json"""    ,	"""wb"""							) as f:
              http_get("""https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json"""    ,	SCREAMING_SNAKE_CASE_							)
           __lowerCamelCase:				int     =	AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""							)
           # The tiny random BERT has a vocab size of 1024, tiny gpt2 as a vocab size of 1000
           self.assertEqual(tokenizer.vocab_size    ,	1000							)
           # Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

        finally:
           os.remove("""tokenizer.json"""							)







     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        __lowerCamelCase:				Optional[int]     =	AlbertTokenizer.from_pretrained("""https://huggingface.co/albert-base-v1/resolve/main/spiece.model"""							)



@is_staging_test
class        a					(			unittest.TestCase   ):
     UpperCAmelCase__ :  Dict						      =	["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]





     @classmethod
     def  SCREAMING_SNAKE_CASE__       (    cls					:				str							):
        __lowerCamelCase:				str     =	TOKEN
        HfFolder.save_token(SCREAMING_SNAKE_CASE_							)





     @classmethod
     def  SCREAMING_SNAKE_CASE__       (    cls					:				List[str]							):







        try:
           delete_repo(token=cls._token    ,	repo_id="""test-tokenizer"""							)
        except HTTPError:
           pass

        try:
           delete_repo(token=cls._token    ,	repo_id="""valid_org/test-tokenizer-org"""							)
        except HTTPError:
           pass

        try:
           delete_repo(token=cls._token    ,	repo_id="""test-dynamic-tokenizer"""							)
        except HTTPError:
           pass





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):







        with tempfile.TemporaryDirectory() as tmp_dir:
           __lowerCamelCase:				str     =	os.path.join(SCREAMING_SNAKE_CASE_    ,	"""vocab.txt"""							)
           with open(SCREAMING_SNAKE_CASE_    ,	"""w"""    ,	encoding="""utf-8"""							) as vocab_writer:
              vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens]							)							)
           __lowerCamelCase:				List[Any]     =	BertTokenizer(SCREAMING_SNAKE_CASE_							)

        tokenizer.push_to_hub("""test-tokenizer"""    ,	use_auth_token=self._token							)
        __lowerCamelCase:				Optional[Any]     =	BertTokenizer.from_pretrained(F'{USER}/test-tokenizer'							)
        self.assertDictEqual(new_tokenizer.vocab    ,	tokenizer.vocab							)

        # Reset repo
        delete_repo(token=self._token    ,	repo_id="""test-tokenizer"""							)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
           tokenizer.save_pretrained(SCREAMING_SNAKE_CASE_    ,	repo_id="""test-tokenizer"""    ,	push_to_hub=SCREAMING_SNAKE_CASE_    ,	use_auth_token=self._token							)

        __lowerCamelCase:				List[Any]     =	BertTokenizer.from_pretrained(F'{USER}/test-tokenizer'							)
        self.assertDictEqual(new_tokenizer.vocab    ,	tokenizer.vocab							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):







        with tempfile.TemporaryDirectory() as tmp_dir:
           __lowerCamelCase:				Union[str, Any]     =	os.path.join(SCREAMING_SNAKE_CASE_    ,	"""vocab.txt"""							)
           with open(SCREAMING_SNAKE_CASE_    ,	"""w"""    ,	encoding="""utf-8"""							) as vocab_writer:
              vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens]							)							)
           __lowerCamelCase:				Optional[Any]     =	BertTokenizer(SCREAMING_SNAKE_CASE_							)

        tokenizer.push_to_hub("""valid_org/test-tokenizer-org"""    ,	use_auth_token=self._token							)
        __lowerCamelCase:				Union[str, Any]     =	BertTokenizer.from_pretrained("""valid_org/test-tokenizer-org"""							)
        self.assertDictEqual(new_tokenizer.vocab    ,	tokenizer.vocab							)

        # Reset repo
        delete_repo(token=self._token    ,	repo_id="""valid_org/test-tokenizer-org"""							)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
           tokenizer.save_pretrained(
               SCREAMING_SNAKE_CASE_    ,	repo_id="""valid_org/test-tokenizer-org"""    ,	push_to_hub=SCREAMING_SNAKE_CASE_    ,	use_auth_token=self._token							)

        __lowerCamelCase:				Dict     =	BertTokenizer.from_pretrained("""valid_org/test-tokenizer-org"""							)
        self.assertDictEqual(new_tokenizer.vocab    ,	tokenizer.vocab							)







     @require_tokenizers
     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        CustomTokenizer.register_for_auto_class()
        with tempfile.TemporaryDirectory() as tmp_dir:
           __lowerCamelCase:				Dict     =	os.path.join(SCREAMING_SNAKE_CASE_    ,	"""vocab.txt"""							)
           with open(SCREAMING_SNAKE_CASE_    ,	"""w"""    ,	encoding="""utf-8"""							) as vocab_writer:
              vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens]							)							)
           __lowerCamelCase:				List[str]     =	CustomTokenizer(SCREAMING_SNAKE_CASE_							)

        # No fast custom tokenizer
        tokenizer.push_to_hub("""test-dynamic-tokenizer"""    ,	use_auth_token=self._token							)

        __lowerCamelCase:				Any     =	AutoTokenizer.from_pretrained(F'{USER}/test-dynamic-tokenizer'    ,	trust_remote_code=SCREAMING_SNAKE_CASE_							)
        # Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
        self.assertEqual(tokenizer.__class__.__name__    ,	"""CustomTokenizer"""							)

        # Fast and slow custom tokenizer
        CustomTokenizerFast.register_for_auto_class()
        with tempfile.TemporaryDirectory() as tmp_dir:
           __lowerCamelCase:				int     =	os.path.join(SCREAMING_SNAKE_CASE_    ,	"""vocab.txt"""							)
           with open(SCREAMING_SNAKE_CASE_    ,	"""w"""    ,	encoding="""utf-8"""							) as vocab_writer:
              vocab_writer.write("""""".join([x + """\n""" for x in self.vocab_tokens]							)							)

           __lowerCamelCase:				str     =	BertTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE_							)
           bert_tokenizer.save_pretrained(SCREAMING_SNAKE_CASE_							)
           __lowerCamelCase:				Union[str, Any]     =	CustomTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE_							)

        tokenizer.push_to_hub("""test-dynamic-tokenizer"""    ,	use_auth_token=self._token							)

        __lowerCamelCase:				Optional[int]     =	AutoTokenizer.from_pretrained(F'{USER}/test-dynamic-tokenizer'    ,	trust_remote_code=SCREAMING_SNAKE_CASE_							)
        # Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
        self.assertEqual(tokenizer.__class__.__name__    ,	"""CustomTokenizerFast"""							)
        __lowerCamelCase:				List[Any]     =	AutoTokenizer.from_pretrained(
            F'{USER}/test-dynamic-tokenizer'    ,	use_fast=SCREAMING_SNAKE_CASE_    ,	trust_remote_code=SCREAMING_SNAKE_CASE_							)
        # Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
        self.assertEqual(tokenizer.__class__.__name__    ,	"""CustomTokenizer"""							)



class        a					(			unittest.TestCase   ):





     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict							):
        __lowerCamelCase:				Union[str, Any]     =	Trie()
        trie.add("""Hello 友達"""							)
        self.assertEqual(trie.data    ,	{"""H""": {"""e""": {"""l""": {"""l""": {"""o""": {""" """: {"""友""": {"""達""": {"""""": 1}}}}}}}}}							)
        trie.add("""Hello"""							)
        trie.data
        self.assertEqual(trie.data    ,	{"""H""": {"""e""": {"""l""": {"""l""": {"""o""": {"""""": 1, """ """: {"""友""": {"""達""": {"""""": 1}}}}}}}}}							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]							):
        __lowerCamelCase:				Optional[Any]     =	Trie()
        self.assertEqual(trie.split("""[CLS] This is a extra_id_100"""							)    ,	["""[CLS] This is a extra_id_100"""]							)
        trie.add("""[CLS]"""							)
        trie.add("""extra_id_1"""							)
        trie.add("""extra_id_100"""							)
        self.assertEqual(trie.split("""[CLS] This is a extra_id_100"""							)    ,	["""[CLS]""", """ This is a """, """extra_id_100"""]							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        __lowerCamelCase:				Union[str, Any]     =	Trie()
        trie.add("""A"""							)
        self.assertEqual(trie.split("""ABC"""							)    ,	["""A""", """BC"""]							)
        self.assertEqual(trie.split("""BCA"""							)    ,	["""BC""", """A"""]							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        __lowerCamelCase:				int     =	Trie()
        trie.add("""TOKEN]"""							)
        trie.add("""[SPECIAL_TOKEN]"""							)
        self.assertEqual(trie.split("""This is something [SPECIAL_TOKEN]"""							)    ,	["""This is something """, """[SPECIAL_TOKEN]"""]							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]							):
        __lowerCamelCase:				Any     =	Trie()
        trie.add("""A"""							)
        trie.add("""P"""							)
        trie.add("""[SPECIAL_TOKEN]"""							)
        self.assertEqual(trie.split("""This is something [SPECIAL_TOKEN]"""							)    ,	["""This is something """, """[SPECIAL_TOKEN]"""]							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        __lowerCamelCase:				str     =	Trie()
        trie.add("""AB"""							)
        trie.add("""B"""							)
        trie.add("""C"""							)
        self.assertEqual(trie.split("""ABC"""							)    ,	["""AB""", """C"""]							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple							):
        __lowerCamelCase:				Any     =	Trie()
        trie.add("""ABC"""							)
        trie.add("""B"""							)
        trie.add("""CD"""							)
        self.assertEqual(trie.split("""ABCD"""							)    ,	["""ABC""", """D"""]							)







     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict							):
        __lowerCamelCase:				Any     =	Trie()
        __lowerCamelCase:				str     =	trie.cut_text("""ABC"""    ,	[0, 0, 2, 1, 2, 3]							)
        self.assertEqual(SCREAMING_SNAKE_CASE_    ,	["""AB""", """C"""]							)






import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map



def     __lowerCAmelCase						(							snake_case	:    Union[dict, list, tuple, torch.Tensor] )       ->							List[Tuple[int, ...]]:
   __lowerCamelCase:				Dict     =	[]
   if isinstance(snake_case						,      snake_case ):
      for v in tree.values():
         shapes.extend(_fetch_dims(snake_case ) )
   elif isinstance(snake_case						,      (list, tuple) ):
      for t in tree:
         shapes.extend(_fetch_dims(snake_case ) )
   elif isinstance(snake_case						,      torch.Tensor ):
      shapes.append(tree.shape )
   else:
      raise ValueError("""Not supported""" )

   return shapes



@torch.jit.ignore
def     __lowerCAmelCase						(							snake_case	:    int						,      snake_case	:    Tuple[int, ...] )       ->							Tuple[int, ...]:
   __lowerCamelCase:				Union[str, Any]     =	[]
   for d in reversed(snake_case ):
      idx.append(flat_idx % d )
      __lowerCamelCase:				Optional[int]     =	flat_idx // d

   return tuple(reversed(snake_case ) )



@torch.jit.ignore
def     __lowerCAmelCase						(							snake_case	:    Sequence[int]						,      snake_case	:    Sequence[int]						,      snake_case	:    Sequence[int]						,      snake_case	:    Optional[Sequence[bool]] = None						,      snake_case	:    Optional[Sequence[bool]] = None						,      )       ->							List[Tuple[slice, ...]]:

   # start_edges and end_edges both indicate whether, starting from any given
   # dimension, the start/end index is at the top/bottom edge of the
   # corresponding tensor, modeled as a tree
   def reduce_edge_list(snake_case	:    List[bool] ) -> None:
      __lowerCamelCase:				List[Any]     =	True
      for i in range(len(snake_case ) ):
         __lowerCamelCase:				str     =	-1 * (i + 1)
         l[reversed_idx] &= tally
         __lowerCamelCase:				List[Any]     =	l[reversed_idx]

   if start_edges is None:
      __lowerCamelCase:				str     =	[s == 0 for s in start]
      reduce_edge_list(snake_case )
   if end_edges is None:
      __lowerCamelCase:				Union[str, Any]     =	[e == (d - 1) for e, d in zip(snake_case						,      snake_case )]
      reduce_edge_list(snake_case )

   # Base cases. Either start/end are empty and we're done, or the final,
   # one-dimensional tensor can be simply sliced
   if len(snake_case ) == 0:
      return [()]
   elif len(snake_case ) == 1:
      return [(slice(start[0]						,      end[0] + 1 ),)]

   __lowerCamelCase:				List[Tuple[slice, ...]]     =	[]
   __lowerCamelCase:				List[slice]     =	[]

   # Dimensions common to start and end can be selected directly
   for s, e in zip(snake_case						,      snake_case ):
      if s == e:
         path_list.append(slice(snake_case						,      s + 1 ) )
      else:
         break

   __lowerCamelCase:				Tuple[slice, ...]     =	tuple(snake_case )
   __lowerCamelCase:				Dict     =	len(snake_case )

   # start == end, and we're done
   if divergence_idx == len(snake_case ):
      return [path]

   def upper() -> Tuple[Tuple[slice, ...], ...]:
      assert start_edges is not None
      assert end_edges is not None

      __lowerCamelCase:				Dict     =	start[divergence_idx]
      return tuple(
          path + (slice(snake_case						,      sdi + 1 ),) + s
          for s in _get_minimal_slice_set(
              start[divergence_idx + 1 :]						,      [d - 1 for d in dims[divergence_idx + 1 :]]						,      dims[divergence_idx + 1 :]						,      start_edges=start_edges[divergence_idx + 1 :]						,      end_edges=[True for _ in end_edges[divergence_idx + 1 :]]						,      ) )

   def lower() -> Tuple[Tuple[slice, ...], ...]:
      assert start_edges is not None
      assert end_edges is not None

      __lowerCamelCase:				List[Any]     =	end[divergence_idx]
      return tuple(
          path + (slice(snake_case						,      edi + 1 ),) + s
          for s in _get_minimal_slice_set(
              [0 for _ in start[divergence_idx + 1 :]]						,      end[divergence_idx + 1 :]						,      dims[divergence_idx + 1 :]						,      start_edges=[True for _ in start_edges[divergence_idx + 1 :]]						,      end_edges=end_edges[divergence_idx + 1 :]						,      ) )

   # If both start and end are at the edges of the subtree rooted at
   # divergence_idx, we can just select the whole subtree at once
   if start_edges[divergence_idx] and end_edges[divergence_idx]:
      slices.append(path + (slice(start[divergence_idx]						,      end[divergence_idx] + 1 ),) )
   # If just start is at the edge, we can grab almost all of the subtree,
   # treating only the ragged bottom edge as an edge case
   elif start_edges[divergence_idx]:
      slices.append(path + (slice(start[divergence_idx]						,      end[divergence_idx] ),) )
      slices.extend(lower() )
   # Analogous to the previous case, but the top is ragged this time
   elif end_edges[divergence_idx]:
      slices.extend(upper() )
      slices.append(path + (slice(start[divergence_idx] + 1						,      end[divergence_idx] + 1 ),) )
   # If both sides of the range are ragged, we need to handle both sides
   # separately. If there's contiguous meat in between them, we can index it
   # in one big chunk
   else:
      slices.extend(upper() )
      __lowerCamelCase:				List[str]     =	end[divergence_idx] - start[divergence_idx]
      if middle_ground > 1:
         slices.append(path + (slice(start[divergence_idx] + 1						,      end[divergence_idx] ),) )
      slices.extend(lower() )

   return slices



@torch.jit.ignore
def     __lowerCAmelCase						(							snake_case	:    torch.Tensor						,      snake_case	:    int						,      snake_case	:    int						,      snake_case	:    int )       ->							torch.Tensor:
   __lowerCamelCase:				Any     =	t.shape[:no_batch_dims]
   __lowerCamelCase:				Dict     =	list(_flat_idx_to_idx(snake_case						,      snake_case ) )
   # _get_minimal_slice_set is inclusive
   __lowerCamelCase:				Optional[int]     =	list(_flat_idx_to_idx(flat_end - 1						,      snake_case ) )

   # Get an ordered list of slices to perform
   __lowerCamelCase:				str     =	_get_minimal_slice_set(
       snake_case						,      snake_case						,      snake_case						,      )

   __lowerCamelCase:				Tuple     =	[t[s] for s in slices]

   return torch.cat([s.view((-1,) + t.shape[no_batch_dims:] ) for s in sliced_tensors] )



def     __lowerCAmelCase						(							snake_case	:    Callable						,      snake_case	:    Dict[str, Any]						,      snake_case	:    int						,      snake_case	:    int						,      snake_case	:    bool = False						,      snake_case	:    Any = None						,      snake_case	:    bool = False						,      )       ->							Any:
   if not (len(snake_case ) > 0):
      raise ValueError("""Must provide at least one input""" )

   __lowerCamelCase:				Optional[int]     =	[shape[:no_batch_dims] for shape in _fetch_dims(snake_case )]
   __lowerCamelCase:				int     =	tuple([max(snake_case ) for s in zip(*snake_case )] )

   def _prep_inputs(snake_case	:    torch.Tensor ) -> torch.Tensor:
      if not low_mem:
         if not sum(t.shape[:no_batch_dims] ) == no_batch_dims:
            __lowerCamelCase:				Optional[Any]     =	t.expand(orig_batch_dims + t.shape[no_batch_dims:] )
         __lowerCamelCase:				Union[str, Any]     =	t.reshape(-1						,      *t.shape[no_batch_dims:] )
      else:
         __lowerCamelCase:				Union[str, Any]     =	t.expand(orig_batch_dims + t.shape[no_batch_dims:] )
      return t

   __lowerCamelCase:				Dict[str, Any]     =	tensor_tree_map(_prep_inputs						,      snake_case )
   __lowerCamelCase:				Optional[Any]     =	None
   if _out is not None:
      __lowerCamelCase:				List[Any]     =	tensor_tree_map(lambda snake_case : t.view([-1] + list(t.shape[no_batch_dims:] ) )						,      _out )

   __lowerCamelCase:				List[Any]     =	1
   for d in orig_batch_dims:
      flat_batch_dim *= d

   __lowerCamelCase:				int     =	flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

   def _select_chunk(snake_case	:    torch.Tensor ) -> torch.Tensor:
      return t[i : i + chunk_size] if t.shape[0] != 1 else t

   __lowerCamelCase:				List[Any]     =	0
   __lowerCamelCase:				Optional[Any]     =	prepped_outputs
   for _ in range(snake_case ):
      # Chunk the input
      if not low_mem:
         __lowerCamelCase:				int     =	_select_chunk
      else:
         __lowerCamelCase:				Optional[int]     =	partial(
             _chunk_slice						,      flat_start=snake_case						,      flat_end=min(snake_case						,      i + chunk_size )						,      no_batch_dims=len(snake_case )						,      )

      __lowerCamelCase:				Dict[str, Any]     =	tensor_tree_map(snake_case						,      snake_case )

      # Run the layer on the chunk
      __lowerCamelCase:				Union[str, Any]     =	layer(**snake_case )

      # Allocate space for the output
      if out is None:
         __lowerCamelCase:				Dict     =	tensor_tree_map(lambda snake_case : t.new_zeros((flat_batch_dim,) + t.shape[1:] )						,      snake_case )

      # Put the chunk in its pre-allocated space
      if isinstance(snake_case						,      snake_case ):

         def assign(snake_case	:    dict						,      snake_case	:    dict ) -> None:
            for k, v in da.items():
               if isinstance(snake_case						,      snake_case ):
                  assign(snake_case						,      da[k] )
               else:
                  if _add_into_out:
                     v[i : i + chunk_size] += da[k]
                  else:
                     __lowerCamelCase:				List[str]     =	da[k]

         assign(snake_case						,      snake_case )
      elif isinstance(snake_case						,      snake_case ):
         for xa, xa in zip(snake_case						,      snake_case ):
            if _add_into_out:
               xa[i : i + chunk_size] += xa
            else:
               __lowerCamelCase:				Optional[int]     =	xa
      elif isinstance(snake_case						,      torch.Tensor ):
         if _add_into_out:
            out[i : i + chunk_size] += output_chunk
         else:
            __lowerCamelCase:				Optional[Any]     =	output_chunk
      else:
         raise ValueError("""Not supported""" )

      i += chunk_size

   __lowerCamelCase:				Union[str, Any]     =	tensor_tree_map(lambda snake_case : t.view(orig_batch_dims + t.shape[1:] )						,      snake_case )

   return out







class        a					:





     def __init__(    self					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				int = 512    ,	):
        __lowerCamelCase:				Any     =	max_chunk_size
        __lowerCamelCase:				Optional[int]     =	None
        __lowerCamelCase:				Optional[tuple]     =	None





     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple    ,	SCREAMING_SNAKE_CASE_					:				Callable    ,	SCREAMING_SNAKE_CASE_					:				tuple    ,	SCREAMING_SNAKE_CASE_					:				int							):
        logging.info("""Tuning chunk size..."""							)

        if min_chunk_size >= self.max_chunk_size:
           return min_chunk_size

        __lowerCamelCase:				List[int]     =	[2**l for l in range(int(math.log(self.max_chunk_size    ,	2							)							) + 1							)]
        __lowerCamelCase:				Optional[Any]     =	[c for c in candidates if c > min_chunk_size]
        __lowerCamelCase:				Union[str, Any]     =	[min_chunk_size] + candidates
        candidates[-1] += 4

        def test_chunk_size(SCREAMING_SNAKE_CASE_					:				int							) -> bool:
           try:
              with torch.no_grad():
                 fn(*SCREAMING_SNAKE_CASE_    ,	chunk_size=SCREAMING_SNAKE_CASE_							)
              return True
           except RuntimeError:
              return False

        __lowerCamelCase:				List[str]     =	0
        __lowerCamelCase:				Any     =	len(SCREAMING_SNAKE_CASE_							) - 1
        while i > min_viable_chunk_size_index:
           __lowerCamelCase:				List[str]     =	test_chunk_size(candidates[i]							)
           if not viable:
              __lowerCamelCase:				Optional[int]     =	(min_viable_chunk_size_index + i) // 2
           else:
              __lowerCamelCase:				Optional[Any]     =	i
              __lowerCamelCase:				List[str]     =	(i + len(SCREAMING_SNAKE_CASE_							) - 1) // 2

        return candidates[min_viable_chunk_size_index]





     def  SCREAMING_SNAKE_CASE__       (    self					:				int    ,	SCREAMING_SNAKE_CASE_					:				Iterable    ,	SCREAMING_SNAKE_CASE_					:				Iterable							):
        __lowerCamelCase:				int     =	True
        for aa, aa in zip(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							):
           assert type(SCREAMING_SNAKE_CASE_							) == type(SCREAMING_SNAKE_CASE_							)
           if isinstance(SCREAMING_SNAKE_CASE_    ,	(list, tuple)							):
              consistent &= self._compare_arg_caches(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)
           elif isinstance(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							):
              __lowerCamelCase:				int     =	[v for _, v in sorted(aa.items()    ,	key=lambda SCREAMING_SNAKE_CASE_							: x[0]							)]
              __lowerCamelCase:				Dict     =	[v for _, v in sorted(aa.items()    ,	key=lambda SCREAMING_SNAKE_CASE_							: x[0]							)]
              consistent &= self._compare_arg_caches(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)
           else:
              consistent &= aa == aa

        return consistent







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]    ,	SCREAMING_SNAKE_CASE_					:				Callable    ,	SCREAMING_SNAKE_CASE_					:				tuple    ,	SCREAMING_SNAKE_CASE_					:				int    ,	):
        __lowerCamelCase:				int     =	True
        __lowerCamelCase:				tuple     =	tree_map(lambda SCREAMING_SNAKE_CASE_							: a.shape if isinstance(SCREAMING_SNAKE_CASE_    ,	torch.Tensor							) else a    ,	SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)
        if self.cached_arg_data is not None:
           # If args have changed shape/value, we need to re-tune
           assert len(self.cached_arg_data							) == len(SCREAMING_SNAKE_CASE_							)
           __lowerCamelCase:				List[str]     =	self._compare_arg_caches(self.cached_arg_data    ,	SCREAMING_SNAKE_CASE_							)
        else:
           # Otherwise, we can reuse the precomputed value
           __lowerCamelCase:				Optional[int]     =	False

        if not consistent:
           __lowerCamelCase:				str     =	self._determine_favorable_chunk_size(
               SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	)
           __lowerCamelCase:				Optional[int]     =	arg_data

        assert self.cached_chunk_size is not None

        return self.cached_chunk_size


