


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   (SCREAMING_SNAKE_CASE )->  List[Tuple[int, ...]]:




      """simple docstring"""

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

      return shapes






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




      """simple docstring"""

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

      return tuple(reversed(SCREAMING_SNAKE_CASE ) )






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




      """simple docstring"""

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

      if start_edges is None:
            snake_case_	=				[s == 0 for s in start]
            reduce_edge_list(SCREAMING_SNAKE_CASE )
      if end_edges is None:
            snake_case_	=				[e == (d - 1) for e, d in zip(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )]
            reduce_edge_list(SCREAMING_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(SCREAMING_SNAKE_CASE ) == 0:
            return [()]
      elif len(SCREAMING_SNAKE_CASE ) == 1:
            return [(slice(start[0] , end[0] + 1 ),)]

      snake_case_	=				[]
      snake_case_	=				[]

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

      snake_case_	=				tuple(SCREAMING_SNAKE_CASE )
      snake_case_	=				len(SCREAMING_SNAKE_CASE )

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

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

            snake_case_	=				start[divergence_idx]
            return tuple(
                path + (slice(SCREAMING_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

            snake_case_	=				end[divergence_idx]
            return tuple(
                path + (slice(SCREAMING_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() )
            snake_case_	=				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   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  torch.Tensor:




      """simple docstring"""

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

      # Get an ordered list of slices to perform
      snake_case_	=				_get_minimal_slice_set(
          SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , )

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

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE = False , SCREAMING_SNAKE_CASE = None , SCREAMING_SNAKE_CASE = False , )->  Any:




      """simple docstring"""

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

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

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

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

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

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

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

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

            snake_case_	=				tensor_tree_map(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

            # Run the layer on the chunk
            snake_case_	=				layer(**SCREAMING_SNAKE_CASE )

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

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

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

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

            i += chunk_size

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

      return out







class lowerCAmelCase_						:

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase = 5_12		,				):
       snake_case_	=				max_chunk_size
       snake_case_	=				None
       snake_case_	=				None



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       logging.info('''Tuning chunk size...''' )

       if min_chunk_size >= self.max_chunk_size:
             return min_chunk_size

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

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

       return candidates[min_viable_chunk_size_index]



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				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 ):
                   snake_case_	=				[v for _, v in sorted(aa.items()		,				key=lambda _UpperCAmelCase : x[0] )]
                   snake_case_	=				[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 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				):
       snake_case_	=				True
       snake_case_	=				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 )
             snake_case_	=				self._compare_arg_caches(self.cached_arg_data		,				_UpperCAmelCase )
       else:
             # Otherwise, we can reuse the precomputed value
             snake_case_	=				False

       if not consistent:
             snake_case_	=				self._determine_favorable_chunk_size(
                 _UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				)
             snake_case_	=				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      =   {"""configuration_vit_msn""": ["""VIT_MSN_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViTMSNConfig"""]}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """ViTMSNModel""",
       """ViTMSNForImageClassification""",
       """ViTMSNPreTrainedModel""",
   ]

if TYPE_CHECKING:
   from .configuration_vit_msn import VIT_MSN_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTMSNConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_vit_msn import (
          VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST,
          ViTMSNForImageClassification,
          ViTMSNModel,
          ViTMSNPreTrainedModel,
      )

else:
   import sys

   UpperCAmelCase      =   _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)