


'''simple docstring'''
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      a (     __a  )    ->      List[Tuple[int, ...]]:



					'''simple docstring'''






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

					return shapes


@torch.jit.ignore
def      a (     __a					,   __a  )    ->      Tuple[int, ...]:



					'''simple docstring'''






					UpperCamelCase__  :Optional[Any]				   =					[]
					for d in reversed(__a  ):
										idx.append(flat_idx % d  )
										UpperCamelCase__  :Union[str, Any]				   =					flat_idx // d

					return tuple(reversed(__a  )  )


@torch.jit.ignore
def      a (     __a					,   __a					,   __a					,   __a = None					,   __a = None					,   )    ->      List[Tuple[slice, ...]]:



					'''simple docstring'''






					def reduce_edge_list(__a  ) -> None:
										UpperCamelCase__  :Union[str, Any]				   =					True
										for i in range(len(__a  )  ):
															UpperCamelCase__  :List[str]				   =					-1 * (i + 1)
															l[reversed_idx] &= tally
															UpperCamelCase__  :Union[str, Any]				   =					l[reversed_idx]

					if start_edges is None:
										UpperCamelCase__  :Optional[Any]				   =					[s == 0 for s in start]
										reduce_edge_list(__a  )
					if end_edges is None:
										UpperCamelCase__  :Tuple				   =					[e == (d - 1) for e, d in zip(__a					,   __a  )]
										reduce_edge_list(__a  )

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

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

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

					UpperCamelCase__  :Tuple[slice, ...]				   =					tuple(__a  )
					UpperCamelCase__  :List[Any]				   =					len(__a  )

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

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

										UpperCamelCase__  :List[Any]				   =					start[divergence_idx]
										return tuple(
										    path + (slice(__a					,   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

										UpperCamelCase__  :int				   =					end[divergence_idx]
										return tuple(
										    path + (slice(__a					,   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()  )
										UpperCamelCase__  :Optional[int]				   =					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      a (     __a					,   __a					,   __a					,   __a  )    ->      torch.Tensor:



					'''simple docstring'''






					UpperCamelCase__  :Any				   =					t.shape[:no_batch_dims]
					UpperCamelCase__  :Tuple				   =					list(_flat_idx_to_idx(__a					,   __a  )  )
					# _get_minimal_slice_set is inclusive
					UpperCamelCase__  :List[Any]				   =					list(_flat_idx_to_idx(flat_end - 1					,   __a  )  )

					# Get an ordered list of slices to perform
					UpperCamelCase__  :int				   =					_get_minimal_slice_set(
					    __a					,   __a					,   __a					,   )

					UpperCamelCase__  :str				   =					[t[s] for s in slices]

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


def      a (     __a					,   __a					,   __a					,   __a					,   __a = False					,   __a = None					,   __a = False					,   )    ->      Any:



					'''simple docstring'''






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

					UpperCamelCase__  :Any				   =					[shape[:no_batch_dims] for shape in _fetch_dims(__a  )]
					UpperCamelCase__  :Tuple				   =					tuple([max(__a  ) for s in zip(*__a  )]  )

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

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

					UpperCamelCase__  :str				   =					1
					for d in orig_batch_dims:
										flat_batch_dim *= d

					UpperCamelCase__  :Any				   =					flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

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

										UpperCamelCase__  :Dict[str, Any]				   =					tensor_tree_map(__a					,   __a  )

										# Run the layer on the chunk
										UpperCamelCase__  :int				   =					layer(**__a  )

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

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

															def assign(__a					,   __a  ) -> None:
																				for k, v in da.items():
																									if isinstance(__a					,   __a  ):
																														assign(__a					,   da[k]  )
																									else:
																														if _add_into_out:
																																			v[i : i + chunk_size] += da[k]
																														else:
																																			UpperCamelCase__  :List[Any]				   =					da[k]

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

										i += chunk_size

					UpperCamelCase__  :List[str]				   =					tensor_tree_map(lambda __a  : t.view(orig_batch_dims + t.shape[1:]  )					,   __a  )

					return out







class   lowercase      :

		"""simple docstring"""

		def __init__(						self       ,  UpperCamelCase_ = 512       ,  ):

							'''simple docstring'''


							UpperCamelCase__  :int				   =					max_chunk_size
							UpperCamelCase__  :Optional[int]				   =					None
							UpperCamelCase__  :Optional[tuple]				   =					None




		def        lowerCAmelCase__							(						self       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_				):

							'''simple docstring'''


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

							if min_chunk_size >= self.max_chunk_size:
												return min_chunk_size

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

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

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

							return candidates[min_viable_chunk_size_index]




		def        lowerCAmelCase__							(						self       ,  UpperCamelCase_       ,  UpperCamelCase_				):

							'''simple docstring'''


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

							return consistent




		def        lowerCAmelCase__							(						self       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  ):

							'''simple docstring'''


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

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

							assert self.cached_chunk_size is not None

							return self.cached_chunk_size



'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__snake_case				 =	{
    '''configuration_convbert''': ['''CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ConvBertConfig''', '''ConvBertOnnxConfig'''],
    '''tokenization_convbert''': ['''ConvBertTokenizer'''],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__snake_case				 =	['''ConvBertTokenizerFast''']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__snake_case				 =	[
		    '''CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''ConvBertForMaskedLM''',
		    '''ConvBertForMultipleChoice''',
		    '''ConvBertForQuestionAnswering''',
		    '''ConvBertForSequenceClassification''',
		    '''ConvBertForTokenClassification''',
		    '''ConvBertLayer''',
		    '''ConvBertModel''',
		    '''ConvBertPreTrainedModel''',
		    '''load_tf_weights_in_convbert''',
		]


try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__snake_case				 =	[
		    '''TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''TFConvBertForMaskedLM''',
		    '''TFConvBertForMultipleChoice''',
		    '''TFConvBertForQuestionAnswering''',
		    '''TFConvBertForSequenceClassification''',
		    '''TFConvBertForTokenClassification''',
		    '''TFConvBertLayer''',
		    '''TFConvBertModel''',
		    '''TFConvBertPreTrainedModel''',
		]


if TYPE_CHECKING:
		from .configuration_convbert import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvBertConfig, ConvBertOnnxConfig
		from .tokenization_convbert import ConvBertTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_convbert_fast import ConvBertTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_convbert import (
				    CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ConvBertForMaskedLM,
				    ConvBertForMultipleChoice,
				    ConvBertForQuestionAnswering,
				    ConvBertForSequenceClassification,
				    ConvBertForTokenClassification,
				    ConvBertLayer,
				    ConvBertModel,
				    ConvBertPreTrainedModel,
				    load_tf_weights_in_convbert,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_convbert import (
				    TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFConvBertForMaskedLM,
				    TFConvBertForMultipleChoice,
				    TFConvBertForQuestionAnswering,
				    TFConvBertForSequenceClassification,
				    TFConvBertForTokenClassification,
				    TFConvBertLayer,
				    TFConvBertModel,
				    TFConvBertPreTrainedModel,
				)


else:
		import sys

		__snake_case				 =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)