






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 				snake_case  (     lowerCamelCase					):




	'''simple docstring'''







	__lowercase             =  []
	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 				snake_case  (     lowerCamelCase				,  lowerCamelCase					):




	'''simple docstring'''







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

	return tuple(reversed(UpperCamelCase__					)					)





@torch.jit.ignore
def 				snake_case  (     lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase = None				,  lowerCamelCase = None				,  ):




	'''simple docstring'''







	def reduce_edge_list(lowerCamelCase					) -> None:
		__lowercase             =  True
		for i in range(len(UpperCamelCase__					)					):
			__lowercase             =  -1 * (i + 1)
			l[reversed_idx] &= tally
			__lowercase             =  l[reversed_idx]

	if start_edges is None:
		__lowercase             =  [s == 0 for s in start]
		reduce_edge_list(UpperCamelCase__					)
	if end_edges is None:
		__lowercase             =  [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					),)]

	__lowercase             =  []
	__lowercase             =  []

	# 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

	__lowercase             =  tuple(UpperCamelCase__					)
	__lowercase             =  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

		__lowercase             =  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

		__lowercase             =  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()					)
		__lowercase             =  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 				snake_case  (     lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase					):




	'''simple docstring'''







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

	# Get an ordered list of slices to perform
	__lowercase             =  _get_minimal_slice_set(
	    UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  )

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

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





def 				snake_case  (     lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase				,  lowerCamelCase = False				,  lowerCamelCase = None				,  lowerCamelCase = False				,  ):




	'''simple docstring'''







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

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

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

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

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

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

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

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

		__lowercase             =  tensor_tree_map(UpperCamelCase__				,  UpperCamelCase__					)

		# Run the layer on the chunk
		__lowercase             =  layer(**UpperCamelCase__					)

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

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

			def assign(lowerCamelCase				,  lowerCamelCase					) -> 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:
							__lowercase             =  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:
					__lowercase             =  xa
		elif isinstance(UpperCamelCase__				,  torch.Tensor					):
			if _add_into_out:
				out[i : i + chunk_size] += output_chunk
			else:
				__lowercase             =  output_chunk
		else:
			raise ValueError("""Not supported"""					)

		i += chunk_size

	__lowercase             =  tensor_tree_map(lambda lowerCamelCase					: t.view(orig_batch_dims + t.shape[1:]					)				,  UpperCamelCase__					)

	return out




class        __UpperCamelCase    :







	def __init__(  self      :							List[str]       ,				_lowerCAmelCase      :							int = 512       ,				)		-> Optional[int]:
		"""simple docstring"""







		__lowercase             =  max_chunk_size
		__lowercase             =  None
		__lowercase             =  None







	def  _a   (  self      :							Union[str, Any]       ,				_lowerCAmelCase      :							Callable       ,				_lowerCAmelCase      :							tuple       ,				_lowerCAmelCase      :							int			)		-> int:
		"""simple docstring"""







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

		if min_chunk_size >= self.max_chunk_size:
			return min_chunk_size

		__lowercase             =  [2**l for l in range(int(math.log(self.max_chunk_size       ,				2			)			) + 1			)]
		__lowercase             =  [c for c in candidates if c > min_chunk_size]
		__lowercase             =  [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

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

		return candidates[min_viable_chunk_size_index]







	def  _a   (  self      :							Dict       ,				_lowerCAmelCase      :							Iterable       ,				_lowerCAmelCase      :							Iterable			)		-> bool:
		"""simple docstring"""







		__lowercase             =  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			):
				__lowercase             =  [v for _, v in sorted(aa.items()       ,				key=lambda _lowerCAmelCase			: x[0]			)]
				__lowercase             =  [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  _a   (  self      :							Tuple       ,				_lowerCAmelCase      :							Callable       ,				_lowerCAmelCase      :							tuple       ,				_lowerCAmelCase      :							int       ,				)		-> int:
		"""simple docstring"""







		__lowercase             =  True
		__lowercase             =  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			)
			__lowercase             =  self._compare_arg_caches(self.cached_arg_data       ,				__a			)
		else:
			# Otherwise, we can reuse the precomputed value
			__lowercase             =  False

		if not consistent:
			__lowercase             =  self._determine_favorable_chunk_size(
			    __a       ,				__a       ,				__a       ,				)
			__lowercase             =  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_torch_available,
)


__UpperCamelCase      :      Tuple     						=						{
    """configuration_swiftformer""": [
        """SWIFTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """SwiftFormerConfig""",
        """SwiftFormerOnnxConfig""",
    ]
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__UpperCamelCase      :      Optional[Any]     						=						[
						    """SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """SwiftFormerForImageClassification""",
						    """SwiftFormerModel""",
						    """SwiftFormerPreTrainedModel""",
						]

if TYPE_CHECKING:
						from .configuration_swiftformer import (
						    SWIFTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    SwiftFormerConfig,
						    SwiftFormerOnnxConfig,
						)

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_swiftformer import (
												    SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    SwiftFormerForImageClassification,
												    SwiftFormerModel,
												    SwiftFormerPreTrainedModel,
												)

else:
						import sys

						__UpperCamelCase      :      List[Any]     						=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

