



'''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__    (     UpperCAmelCase_							):
						_UpperCamelCase					:		Dict 				=   []
						if isinstance(UpperCAmelCase_  ,  UpperCAmelCase_							):
												for v in tree.values():
																		shapes.extend(_fetch_dims(UpperCAmelCase_							)							)
						elif isinstance(UpperCAmelCase_  ,  (list, tuple)							):
												for t in tree:
																		shapes.extend(_fetch_dims(UpperCAmelCase_							)							)
						elif isinstance(UpperCAmelCase_  ,  torch.Tensor							):
												shapes.append(tree.shape							)
						else:
												raise ValueError('Not supported'							)

						return shapes






@torch.jit.ignore
def A__    (     UpperCAmelCase_  ,  UpperCAmelCase_							):
						_UpperCamelCase					:		Optional[Any] 				=   []
						for d in reversed(UpperCAmelCase_							):
												idx.append(flat_idx % d							)
												_UpperCamelCase					:		Optional[Any] 				=   flat_idx // d

						return tuple(reversed(UpperCAmelCase_							)							)






@torch.jit.ignore
def A__    (     UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_ = None  ,  UpperCAmelCase_ = None  ,  ):



						def reduce_edge_list(UpperCAmelCase_							) -> None:
												_UpperCamelCase					:		Optional[int] 				=   True
												for i in range(len(UpperCAmelCase_							)							):
																		_UpperCamelCase					:		Optional[Any] 				=   -1 * (i + 1)
																		l[reversed_idx] &= tally
																		_UpperCamelCase					:		Dict 				=   l[reversed_idx]

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

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

						_UpperCamelCase					:		Optional[int] 				=   []
						_UpperCamelCase					:		Any 				=   []

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

						_UpperCamelCase					:		Optional[int] 				=   tuple(UpperCAmelCase_							)
						_UpperCamelCase					:		Optional[Any] 				=   len(UpperCAmelCase_							)

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

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

												_UpperCamelCase					:		str 				=   start[divergence_idx]
												return tuple(
												    path + (slice(UpperCAmelCase_  ,  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					:		Optional[Any] 				=   end[divergence_idx]
												return tuple(
												    path + (slice(UpperCAmelCase_  ,  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[Any] 				=   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__    (     UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_							):
						_UpperCamelCase					:		Optional[int] 				=   t.shape[:no_batch_dims]
						_UpperCamelCase					:		Optional[Any] 				=   list(_flat_idx_to_idx(UpperCAmelCase_  ,  UpperCAmelCase_							)							)
						# _get_minimal_slice_set is inclusive
						_UpperCamelCase					:		Optional[int] 				=   list(_flat_idx_to_idx(flat_end - 1  ,  UpperCAmelCase_							)							)

						# Get an ordered list of slices to perform
						_UpperCamelCase					:		Tuple 				=   _get_minimal_slice_set(
						    UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_  ,  )

						_UpperCamelCase					:		Any 				=   [t[s] for s in slices]

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






def A__    (     UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_ = False  ,  UpperCAmelCase_ = None  ,  UpperCAmelCase_ = False  ,  ):



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

						_UpperCamelCase					:		Dict 				=   [shape[:no_batch_dims] for shape in _fetch_dims(UpperCAmelCase_							)]
						_UpperCamelCase					:		List[Any] 				=   tuple([max(UpperCAmelCase_							) for s in zip(*UpperCAmelCase_							)]							)

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

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

						_UpperCamelCase					:		Dict 				=   1
						for d in orig_batch_dims:
												flat_batch_dim *= d

						_UpperCamelCase					:		List[str] 				=   flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

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

												_UpperCamelCase					:		Optional[int] 				=   tensor_tree_map(UpperCAmelCase_  ,  UpperCAmelCase_							)

												# Run the layer on the chunk
												_UpperCamelCase					:		List[str] 				=   layer(**UpperCAmelCase_							)

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

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

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

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

												i += chunk_size

						_UpperCamelCase					:		Any 				=   tensor_tree_map(lambda UpperCAmelCase_							: t.view(orig_batch_dims + t.shape[1:]							)  ,  UpperCAmelCase_							)

						return out






class 							lowercase__	:





	def __init__( self					:				Optional[Any] ,lowerCamelCase__					:				int = 512 ,):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[int] 				=   max_chunk_size
							_UpperCamelCase					:		Any 				=   None
							_UpperCamelCase					:		Optional[Any] 				=   None





	def 	UpperCamelCase_  ( self					:				int ,lowerCamelCase__					:				Callable ,lowerCamelCase__					:				tuple ,lowerCamelCase__					:				int ):







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

							if min_chunk_size >= self.max_chunk_size:
													return min_chunk_size

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

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

							_UpperCamelCase					:		Tuple 				=   0
							_UpperCamelCase					:		Dict 				=   len(A_ ) - 1
							while i > min_viable_chunk_size_index:
													_UpperCamelCase					:		Union[str, Any] 				=   test_chunk_size(candidates[i] )
													if not viable:
																			_UpperCamelCase					:		Tuple 				=   (min_viable_chunk_size_index + i) // 2
													else:
																			_UpperCamelCase					:		List[Any] 				=   i
																			_UpperCamelCase					:		List[Any] 				=   (i + len(A_ ) - 1) // 2

							return candidates[min_viable_chunk_size_index]





	def 	UpperCamelCase_  ( self					:				List[str] ,lowerCamelCase__					:				Iterable ,lowerCamelCase__					:				Iterable ):







							'''simple docstring'''
							_UpperCamelCase					:		Any 				=   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_ ):
																			_UpperCamelCase					:		Dict 				=   [v for _, v in sorted(aa.items() ,key=lambda lowerCamelCase__ : x[0] )]
																			_UpperCamelCase					:		int 				=   [v for _, v in sorted(aa.items() ,key=lambda lowerCamelCase__ : x[0] )]
																			consistent &= self._compare_arg_caches(A_ ,A_ )
													else:
																			consistent &= aa == aa

							return consistent





	def 	UpperCamelCase_  ( self					:				int ,lowerCamelCase__					:				Callable ,lowerCamelCase__					:				tuple ,lowerCamelCase__					:				int ,):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[Any] 				=   True
							_UpperCamelCase					:		Tuple 				=   tree_map(lambda lowerCamelCase__ : 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_ )
													_UpperCamelCase					:		int 				=   self._compare_arg_caches(self.cached_arg_data ,A_ )
							else:
													# Otherwise, we can reuse the precomputed value
													_UpperCamelCase					:		int 				=   False

							if not consistent:
													_UpperCamelCase					:		int 				=   self._determine_favorable_chunk_size(
													    A_ ,A_ ,A_ ,)
													_UpperCamelCase					:		Optional[int] 				=   arg_data

							assert self.cached_chunk_size is not None

							return self.cached_chunk_size








from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCamelCase				:		str     =   {
    "configuration_nezha": ["NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP", "NezhaConfig"],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase				:		List[str]     =   [
     "NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST",
     "NezhaForNextSentencePrediction",
     "NezhaForMaskedLM",
     "NezhaForPreTraining",
     "NezhaForMultipleChoice",
     "NezhaForQuestionAnswering",
     "NezhaForSequenceClassification",
     "NezhaForTokenClassification",
     "NezhaModel",
     "NezhaPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_nezha import NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP, NezhaConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nezha import (
      NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST,
      NezhaForMaskedLM,
      NezhaForMultipleChoice,
      NezhaForNextSentencePrediction,
      NezhaForPreTraining,
      NezhaForQuestionAnswering,
      NezhaForSequenceClassification,
      NezhaForTokenClassification,
      NezhaModel,
      NezhaPreTrainedModel,
  )


else:
 import sys

 lowerCamelCase				:		str     =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



