



"""simple docstring"""






import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


_lowercase : List[Any]      						=						logging.get_logger(__name__)

_lowercase : int      						=						"▁"

_lowercase : int      						=						{"vocab_file": "sentencepiece.bpe.model"}

_lowercase : Optional[Any]      						=						{
    "vocab_file": {
        "facebook/nllb-200-distilled-600M": (
            "https://huggingface.co/facebook/nllb-200-distilled-600M/blob/main/sentencepiece.bpe.model"
        ),
    }
}

_lowercase : Dict      						=						{
    "facebook/nllb-200-distilled-600M": 1_0_2_4,
}

# fmt: off
_lowercase : List[str]      						=						["ace_Arab", "ace_Latn", "acm_Arab", "acq_Arab", "aeb_Arab", "afr_Latn", "ajp_Arab", "aka_Latn", "amh_Ethi", "apc_Arab", "arb_Arab", "ars_Arab", "ary_Arab", "arz_Arab", "asm_Beng", "ast_Latn", "awa_Deva", "ayr_Latn", "azb_Arab", "azj_Latn", "bak_Cyrl", "bam_Latn", "ban_Latn", "bel_Cyrl", "bem_Latn", "ben_Beng", "bho_Deva", "bjn_Arab", "bjn_Latn", "bod_Tibt", "bos_Latn", "bug_Latn", "bul_Cyrl", "cat_Latn", "ceb_Latn", "ces_Latn", "cjk_Latn", "ckb_Arab", "crh_Latn", "cym_Latn", "dan_Latn", "deu_Latn", "dik_Latn", "dyu_Latn", "dzo_Tibt", "ell_Grek", "eng_Latn", "epo_Latn", "est_Latn", "eus_Latn", "ewe_Latn", "fao_Latn", "pes_Arab", "fij_Latn", "fin_Latn", "fon_Latn", "fra_Latn", "fur_Latn", "fuv_Latn", "gla_Latn", "gle_Latn", "glg_Latn", "grn_Latn", "guj_Gujr", "hat_Latn", "hau_Latn", "heb_Hebr", "hin_Deva", "hne_Deva", "hrv_Latn", "hun_Latn", "hye_Armn", "ibo_Latn", "ilo_Latn", "ind_Latn", "isl_Latn", "ita_Latn", "jav_Latn", "jpn_Jpan", "kab_Latn", "kac_Latn", "kam_Latn", "kan_Knda", "kas_Arab", "kas_Deva", "kat_Geor", "knc_Arab", "knc_Latn", "kaz_Cyrl", "kbp_Latn", "kea_Latn", "khm_Khmr", "kik_Latn", "kin_Latn", "kir_Cyrl", "kmb_Latn", "kon_Latn", "kor_Hang", "kmr_Latn", "lao_Laoo", "lvs_Latn", "lij_Latn", "lim_Latn", "lin_Latn", "lit_Latn", "lmo_Latn", "ltg_Latn", "ltz_Latn", "lua_Latn", "lug_Latn", "luo_Latn", "lus_Latn", "mag_Deva", "mai_Deva", "mal_Mlym", "mar_Deva", "min_Latn", "mkd_Cyrl", "plt_Latn", "mlt_Latn", "mni_Beng", "khk_Cyrl", "mos_Latn", "mri_Latn", "zsm_Latn", "mya_Mymr", "nld_Latn", "nno_Latn", "nob_Latn", "npi_Deva", "nso_Latn", "nus_Latn", "nya_Latn", "oci_Latn", "gaz_Latn", "ory_Orya", "pag_Latn", "pan_Guru", "pap_Latn", "pol_Latn", "por_Latn", "prs_Arab", "pbt_Arab", "quy_Latn", "ron_Latn", "run_Latn", "rus_Cyrl", "sag_Latn", "san_Deva", "sat_Beng", "scn_Latn", "shn_Mymr", "sin_Sinh", "slk_Latn", "slv_Latn", "smo_Latn", "sna_Latn", "snd_Arab", "som_Latn", "sot_Latn", "spa_Latn", "als_Latn", "srd_Latn", "srp_Cyrl", "ssw_Latn", "sun_Latn", "swe_Latn", "swh_Latn", "szl_Latn", "tam_Taml", "tat_Cyrl", "tel_Telu", "tgk_Cyrl", "tgl_Latn", "tha_Thai", "tir_Ethi", "taq_Latn", "taq_Tfng", "tpi_Latn", "tsn_Latn", "tso_Latn", "tuk_Latn", "tum_Latn", "tur_Latn", "twi_Latn", "tzm_Tfng", "uig_Arab", "ukr_Cyrl", "umb_Latn", "urd_Arab", "uzn_Latn", "vec_Latn", "vie_Latn", "war_Latn", "wol_Latn", "xho_Latn", "ydd_Hebr", "yor_Latn", "yue_Hant", "zho_Hans", "zho_Hant", "zul_Latn"]

class        __SCREAMING_SNAKE_CASE	(      lowerCAmelCase_   ):


					'''simple docstring'''





					_a   			=							VOCAB_FILES_NAMES
					_a   			=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					_a   			=							PRETRAINED_VOCAB_FILES_MAP
					_a   			=							['input_ids', 'attention_mask']

					_a   			=							[]
					_a   			=							[]



					def __init__(       self							:  Optional[int],  lowerCamelCase							:  str,  lowerCamelCase							:  str="<s>",  lowerCamelCase							:  str="</s>",  lowerCamelCase							:  Dict="</s>",  lowerCamelCase							:  List[Any]="<s>",  lowerCamelCase							:  Any="<unk>",  lowerCamelCase							:  List[str]="<pad>",  lowerCamelCase							:  Dict="<mask>",  lowerCamelCase							:  List[str]=None,  lowerCamelCase							:  Optional[int]=None,  lowerCamelCase							:  Tuple=None,  lowerCamelCase							:  Optional[Dict[str, Any]] = None,  lowerCamelCase							:  str=None,  lowerCamelCase							:  str=False,  **lowerCamelCase							:  Any,  )->					Union[str, Any]:
								lowerCamelCase__		:		str			=AddedToken(A_,  lstrip=A_,  rstrip=A_     ) if isinstance(A_,  A_     ) else mask_token

								lowerCamelCase__		:		Optional[Any]			={} if sp_model_kwargs is None else sp_model_kwargs
								lowerCamelCase__		:		Union[str, Any]			=legacy_behaviour

								super().__init__(
								    bos_token=A_,  eos_token=A_,  unk_token=A_,  sep_token=A_,  cls_token=A_,  pad_token=A_,  mask_token=A_,  tokenizer_file=A_,  src_lang=A_,  tgt_lang=A_,  additional_special_tokens=A_,  sp_model_kwargs=self.sp_model_kwargs,  legacy_behaviour=A_,  **A_,  )

								lowerCamelCase__		:		Any			=spm.SentencePieceProcessor(**self.sp_model_kwargs     )
								self.sp_model.Load(str(A_     )     )
								lowerCamelCase__		:		Optional[Any]			=vocab_file
								# Original fairseq vocab and spm vocab must be "aligned":
								# Vocab    |    0    |    1    |   2    |    3    |  4   |  5   |  6   |   7  |   8  |  9
								# -------- | ------- | ------- | ------ | ------- | ---- | ---- | ---- | ---- | ---- | ----
								# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | 'an' | '▁n' | '▁m' | '▁t' | '▁k' | '▁a'
								# spm      | '<unk>' | '<s>'   | '</s>' | 'an'    | '▁n' | '▁m' | '▁t' | '▁k' | '▁a' | '▁s'

								# Mimic fairseq token-to-id alignment for the first 4 token
								lowerCamelCase__		:		Dict			={'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

								# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
								lowerCamelCase__		:		str			=1

								lowerCamelCase__		:		Union[str, Any]			=len(self.sp_model     )
								lowerCamelCase__		:		Any			={
								    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(A_     )
								}
								lowerCamelCase__		:		Optional[Any]			={v: k for k, v in self.lang_code_to_id.items()}
								lowerCamelCase__		:		Tuple			=len(self.sp_model     ) + len(self.lang_code_to_id     ) + self.fairseq_offset

								self.fairseq_tokens_to_ids.update(self.lang_code_to_id     )
								lowerCamelCase__		:		List[Any]			={v: k for k, v in self.fairseq_tokens_to_ids.items()}
								lowerCamelCase__		:		Tuple			=list(self.lang_code_to_id.keys()     )

								if additional_special_tokens is not None:
											# Only add those special tokens if they are not already there.
											self._additional_special_tokens.extend(
											    [t for t in additional_special_tokens if t not in self._additional_special_tokens]     )

								lowerCamelCase__		:		List[Any]			=src_lang if src_lang is not None else '''eng_Latn'''
								lowerCamelCase__		:		Dict			=self.lang_code_to_id[self._src_lang]
								lowerCamelCase__		:		Optional[int]			=tgt_lang
								self.set_src_lang_special_tokens(self._src_lang     )



					def __getstate__(       self							:  Union[str, Any]     )->					Optional[Any]:
								lowerCamelCase__		:		Union[str, Any]			=self.__dict__.copy()
								lowerCamelCase__		:		List[str]			=None
								lowerCamelCase__		:		Any			=self.sp_model.serialized_model_proto()
								return state



					def __setstate__(       self							:  Union[str, Any],  lowerCamelCase							:  List[Any]     )->					List[Any]:
								lowerCamelCase__		:		Any			=d

								# for backward compatibility
								if not hasattr(self,  '''sp_model_kwargs'''     ):
											lowerCamelCase__		:		Optional[Any]			={}

								lowerCamelCase__		:		List[str]			=spm.SentencePieceProcessor(**self.sp_model_kwargs     )
								self.sp_model.LoadFromSerializedProto(self.sp_model_proto     )



					@property
					def 		snake_case       (       self							:  Union[str, Any]     )->					Any:
								return len(self.sp_model     ) + len(self.lang_code_to_id     ) + self.fairseq_offset + 1  # Plus 1 for the mask token



					@property
					def 		snake_case       (       self							:  Any     )->					str:
								return self._src_lang



					@src_lang.setter
					def 		snake_case       (       self							:  Union[str, Any],  lowerCamelCase							:  str     )->					None:
								lowerCamelCase__		:		Any			=new_src_lang
								self.set_src_lang_special_tokens(self._src_lang     )



					def 		snake_case       (       self							:  Any,  lowerCamelCase							:  List[int],  lowerCamelCase							:  Optional[List[int]] = None,  lowerCamelCase							:  bool = False     )->					List[int]:






								if already_has_special_tokens:
											return super().get_special_tokens_mask(
											    token_ids_a=A_,  token_ids_a=A_,  already_has_special_tokens=A_     )

								lowerCamelCase__		:		int			=[1] * len(self.prefix_tokens     )
								lowerCamelCase__		:		Dict			=[1] * len(self.suffix_tokens     )
								if token_ids_a is None:
											return prefix_ones + ([0] * len(A_     )) + suffix_ones
								return prefix_ones + ([0] * len(A_     )) + ([0] * len(A_     )) + suffix_ones



					def 		snake_case       (       self							:  Union[str, Any],  lowerCamelCase							:  List[int],  lowerCamelCase							:  Optional[List[int]] = None     )->					List[int]:






								if token_ids_a is None:
											return self.prefix_tokens + token_ids_a + self.suffix_tokens
								# We don't expect to process pairs, but leave the pair logic for API consistency
								return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens



					def 		snake_case       (       self							:  str,  lowerCamelCase							:  List[int],  lowerCamelCase							:  Optional[List[int]] = None     )->					List[int]:
								lowerCamelCase__		:		Tuple			=[self.sep_token_id]
								lowerCamelCase__		:		List[str]			=[self.cls_token_id]

								if token_ids_a is None:
											return len(cls + token_ids_a + sep     ) * [0]
								return len(cls + token_ids_a + sep + sep + token_ids_a + sep     ) * [0]



					def 		snake_case       (       self							:  Dict,  lowerCamelCase							:  int,  lowerCamelCase							:  str,  lowerCamelCase							:  Optional[str],  lowerCamelCase							:  Optional[str],  **lowerCamelCase							:  str     )->					List[str]:






								if src_lang is None or tgt_lang is None:
											raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''     )
								lowerCamelCase__		:		List[str]			=src_lang
								lowerCamelCase__		:		Any			=self(A_,  add_special_tokens=A_,  return_tensors=A_,  **A_     )
								lowerCamelCase__		:		Optional[Any]			=self.convert_tokens_to_ids(A_     )
								lowerCamelCase__		:		Optional[int]			=tgt_lang_id
								return inputs



					def 		snake_case       (       self							:  Dict     )->					List[Any]:
								lowerCamelCase__		:		Union[str, Any]			={self.convert_ids_to_tokens(A_     ): i for i in range(self.vocab_size     )}
								vocab.update(self.added_tokens_encoder     )
								return vocab



					def 		snake_case       (       self							:  Optional[Any],  lowerCamelCase							:  str     )->					List[str]:
								return self.sp_model.encode(A_,  out_type=A_     )



					def 		snake_case       (       self							:  Optional[int],  lowerCamelCase							:  Optional[int]     )->					Union[str, Any]:






								if token in self.fairseq_tokens_to_ids:
											return self.fairseq_tokens_to_ids[token]
								lowerCamelCase__		:		str			=self.sp_model.PieceToId(A_     )

								# Need to return unknown token if the SP model returned 0
								return spm_id + self.fairseq_offset if spm_id else self.unk_token_id



					def 		snake_case       (       self							:  List[str],  lowerCamelCase							:  List[Any]     )->					List[str]:






								if index in self.fairseq_ids_to_tokens:
											return self.fairseq_ids_to_tokens[index]
								return self.sp_model.IdToPiece(index - self.fairseq_offset     )



					def 		snake_case       (       self							:  Optional[int],  lowerCamelCase							:  str     )->					List[str]:
								lowerCamelCase__		:		Union[str, Any]			=''''''.join(A_     ).replace(A_,  ''' '''     ).strip()
								return out_string



					def 		snake_case       (       self							:  Optional[int],  lowerCamelCase							:  str,  lowerCamelCase							:  Optional[str] = None     )->					Tuple[str]:






								if not os.path.isdir(A_     ):
											logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''     )
											return
								lowerCamelCase__		:		List[Any]			=os.path.join(
								    A_,  (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']     )

								if os.path.abspath(self.vocab_file     ) != os.path.abspath(A_     ) and os.path.isfile(self.vocab_file     ):
											copyfile(self.vocab_file,  A_     )
								elif not os.path.isfile(self.vocab_file     ):
											with open(A_,  '''wb'''     ) as fi:
														lowerCamelCase__		:		str			=self.sp_model.serialized_model_proto()
														fi.write(A_     )

								return (out_vocab_file,)



					def 		snake_case       (       self							:  Dict,  lowerCamelCase							:  List[str],  lowerCamelCase							:  str = "eng_Latn",  lowerCamelCase							:  Optional[List[str]] = None,  lowerCamelCase							:  str = "fra_Latn",  **lowerCamelCase							:  Optional[int],  )->					BatchEncoding:
								lowerCamelCase__		:		Optional[Any]			=src_lang
								lowerCamelCase__		:		List[str]			=tgt_lang
								return super().prepare_seqaseq_batch(A_,  A_,  **A_     )



					def 		snake_case       (       self							:  int     )->					Optional[Any]:
								return self.set_src_lang_special_tokens(self.src_lang     )



					def 		snake_case       (       self							:  Optional[int]     )->					Union[str, Any]:
								return self.set_tgt_lang_special_tokens(self.tgt_lang     )



					def 		snake_case       (       self							:  int,  lowerCamelCase							:  Dict     )->					None:
								lowerCamelCase__		:		Optional[int]			=self.lang_code_to_id[src_lang]
								if self.legacy_behaviour:
											lowerCamelCase__		:		str			=[]
											lowerCamelCase__		:		Union[str, Any]			=[self.eos_token_id, self.cur_lang_code]
								else:
											lowerCamelCase__		:		Tuple			=[self.cur_lang_code]
											lowerCamelCase__		:		List[Any]			=[self.eos_token_id]




					def 		snake_case       (       self							:  List[Any],  lowerCamelCase							:  str     )->					None:
								lowerCamelCase__		:		str			=self.lang_code_to_id[lang]
								if self.legacy_behaviour:
											lowerCamelCase__		:		List[str]			=[]
											lowerCamelCase__		:		str			=[self.eos_token_id, self.cur_lang_code]
								else:
											lowerCamelCase__		:		Dict			=[self.cur_lang_code]
											lowerCamelCase__		:		int			=[self.eos_token_id]










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 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Union[dict, list, tuple, torch.Tensor]     ):


      '''simple docstring'''



      lowerCamelCase_		     =    []
      if isinstance(lowercase	,							lowercase     ):
            for v in tree.values():
                  shapes.extend(_fetch_dims(lowercase     )     )
      elif isinstance(lowercase	,							(list, tuple)     ):
            for t in tree:
                  shapes.extend(_fetch_dims(lowercase     )     )
      elif isinstance(lowercase	,							torch.Tensor     ):
            shapes.append(tree.shape     )
      else:
            raise ValueError('Not supported'     )

      return shapes

@torch.jit.ignore
def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       int	,							lowercase	:       Tuple[int, ...]     ):


      '''simple docstring'''



      lowerCamelCase_		     =    []
      for d in reversed(lowercase     ):
            idx.append(flat_idx % d     )
            lowerCamelCase_		     =    flat_idx // d

      return tuple(reversed(lowercase     )     )

@torch.jit.ignore
def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Sequence[int]	,							lowercase	:       Sequence[int]	,							lowercase	:       Sequence[int]	,							lowercase	:       Optional[Sequence[bool]] = None	,							lowercase	:       Optional[Sequence[bool]] = None	,							):


      '''simple docstring'''



      def reduce_edge_list(lowercase	:       List[bool]     ) -> None:
            lowerCamelCase_		     =    True
            for i in range(len(lowercase     )     ):
                  lowerCamelCase_		     =    -1 * (i + 1)
                  l[reversed_idx] &= tally
                  lowerCamelCase_		     =    l[reversed_idx]

      if start_edges is None:
            lowerCamelCase_		     =    [s == 0 for s in start]
            reduce_edge_list(lowercase     )
      if end_edges is None:
            lowerCamelCase_		     =    [e == (d - 1) for e, d in zip(lowercase	,							lowercase     )]
            reduce_edge_list(lowercase     )

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

      lowerCamelCase_		     =    []
      lowerCamelCase_		     =    []

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

      lowerCamelCase_		     =    tuple(lowercase     )
      lowerCamelCase_		     =    len(lowercase     )

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

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

            lowerCamelCase_		     =    start[divergence_idx]
            return tuple(
                path + (slice(lowercase	,							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_		     =    end[divergence_idx]
            return tuple(
                path + (slice(lowercase	,							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_		     =    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 			_SCREAMING_SNAKE_CASE			(							lowercase	:       torch.Tensor	,							lowercase	:       int	,							lowercase	:       int	,							lowercase	:       int     ):


      '''simple docstring'''



      lowerCamelCase_		     =    t.shape[:no_batch_dims]
      lowerCamelCase_		     =    list(_flat_idx_to_idx(lowercase	,							lowercase     )     )
      # _get_minimal_slice_set is inclusive
      lowerCamelCase_		     =    list(_flat_idx_to_idx(flat_end - 1	,							lowercase     )     )

      # Get an ordered list of slices to perform
      lowerCamelCase_		     =    _get_minimal_slice_set(
          lowercase	,							lowercase	,							lowercase	,							)

      lowerCamelCase_		     =    [t[s] for s in slices]

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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Callable	,							lowercase	:       Dict[str, Any]	,							lowercase	:       int	,							lowercase	:       int	,							lowercase	:       bool = False	,							lowercase	:       Any = None	,							lowercase	:       bool = False	,							):


      '''simple docstring'''



      if not (len(lowercase     ) > 0):
            raise ValueError('Must provide at least one input'     )

      lowerCamelCase_		     =    [shape[:no_batch_dims] for shape in _fetch_dims(lowercase     )]
      lowerCamelCase_		     =    tuple([max(lowercase     ) for s in zip(*lowercase     )]     )

      def _prep_inputs(lowercase	:       torch.Tensor     ) -> torch.Tensor:
            if not low_mem:
                  if not sum(t.shape[:no_batch_dims]     ) == no_batch_dims:
                        lowerCamelCase_		     =    t.expand(orig_batch_dims + t.shape[no_batch_dims:]     )
                  lowerCamelCase_		     =    t.reshape(-1	,							*t.shape[no_batch_dims:]     )
            else:
                  lowerCamelCase_		     =    t.expand(orig_batch_dims + t.shape[no_batch_dims:]     )
            return t

      lowerCamelCase_		     =    tensor_tree_map(_prep_inputs	,							lowercase     )
      lowerCamelCase_		     =    None
      if _out is not None:
            lowerCamelCase_		     =    tensor_tree_map(lambda lowercase     : t.view([-1] + list(t.shape[no_batch_dims:]     )     )	,							_out     )

      lowerCamelCase_		     =    1
      for d in orig_batch_dims:
            flat_batch_dim *= d

      lowerCamelCase_		     =    flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

      lowerCamelCase_		     =    0
      lowerCamelCase_		     =    prepped_outputs
      for _ in range(lowercase     ):
            # Chunk the input
            if not low_mem:
                  lowerCamelCase_		     =    _select_chunk
            else:
                  lowerCamelCase_		     =    partial(
                      _chunk_slice	,							flat_start=lowercase	,							flat_end=min(lowercase	,							i + chunk_size     )	,							no_batch_dims=len(lowercase     )	,							)

            lowerCamelCase_		     =    tensor_tree_map(lowercase	,							lowercase     )

            # Run the layer on the chunk
            lowerCamelCase_		     =    layer(**lowercase     )

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

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

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

                  assign(lowercase	,							lowercase     )
            elif isinstance(lowercase	,							lowercase     ):
                  for xa, xa in zip(lowercase	,							lowercase     ):
                        if _add_into_out:
                              xa[i : i + chunk_size] += xa
                        else:
                              lowerCamelCase_		     =    xa
            elif isinstance(lowercase	,							torch.Tensor     ):
                  if _add_into_out:
                        out[i : i + chunk_size] += output_chunk
                  else:
                        lowerCamelCase_		     =    output_chunk
            else:
                  raise ValueError('Not supported'     )

            i += chunk_size

      lowerCamelCase_		     =    tensor_tree_map(lambda lowercase     : t.view(orig_batch_dims + t.shape[1:]     )	,							lowercase     )

      return out







class   A:



 '''simple docstring'''


 def __init__(				self					:  Optional[Any]						,							A_					:  int = 512						,							)   ->	Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_		     =    max_chunk_size
       lowerCamelCase_		     =    None
       lowerCamelCase_		     =    None

 def 				a__ (				self					:  int						,							A_					:  Callable						,							A_					:  tuple						,							A_					:  int  )   ->	int:

       """simple docstring"""






       logging.info('Tuning chunk size...'  )

       if min_chunk_size >= self.max_chunk_size:
             return min_chunk_size

       lowerCamelCase_		     =    [2**l for l in range(int(math.log(self.max_chunk_size						,							2  )  ) + 1  )]
       lowerCamelCase_		     =    [c for c in candidates if c > min_chunk_size]
       lowerCamelCase_		     =    [min_chunk_size] + candidates
       candidates[-1] += 4

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

       lowerCamelCase_		     =    0
       lowerCamelCase_		     =    len(A_  ) - 1
       while i > min_viable_chunk_size_index:
             lowerCamelCase_		     =    test_chunk_size(candidates[i]  )
             if not viable:
                   lowerCamelCase_		     =    (min_viable_chunk_size_index + i) // 2
             else:
                   lowerCamelCase_		     =    i
                   lowerCamelCase_		     =    (i + len(A_  ) - 1) // 2

       return candidates[min_viable_chunk_size_index]

 def 				a__ (				self					:  List[str]						,							A_					:  Iterable						,							A_					:  Iterable  )   ->	bool:

       """simple docstring"""






       lowerCamelCase_		     =    True
       for aa, aa in zip(A_						,							A_  ):
             assert type(A_  ) == type(A_  )
             if isinstance(A_						,							(list, tuple)  ):
                   consistent &= self._compare_arg_caches(A_						,							A_  )
             elif isinstance(A_						,							A_  ):
                   lowerCamelCase_		     =    [v for _, v in sorted(aa.items()						,							key=lambda A_  : x[0]  )]
                   lowerCamelCase_		     =    [v for _, v in sorted(aa.items()						,							key=lambda A_  : x[0]  )]
                   consistent &= self._compare_arg_caches(A_						,							A_  )
             else:
                   consistent &= aa == aa

       return consistent





 def 				a__ (				self					:  int						,							A_					:  Callable						,							A_					:  tuple						,							A_					:  int						,							)   ->	int:

       """simple docstring"""






       lowerCamelCase_		     =    True
       lowerCamelCase_		     =    tree_map(lambda A_  : a.shape if isinstance(A_						,							torch.Tensor  ) else a						,							A_						,							A_  )
       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(A_  )
             lowerCamelCase_		     =    self._compare_arg_caches(self.cached_arg_data						,							A_  )
       else:
             # Otherwise, we can reuse the precomputed value
             lowerCamelCase_		     =    False

       if not consistent:
             lowerCamelCase_		     =    self._determine_favorable_chunk_size(
                 A_						,							A_						,							A_						,							)
             lowerCamelCase_		     =    arg_data

       assert self.cached_chunk_size is not None

       return self.cached_chunk_size



