import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
      MODEL_MAPPING,
      LevitForImageClassification,
      LevitForImageClassificationWithTeacher,
      LevitModel,
  )
  from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import LevitImageProcessor







class   __lowerCAmelCase			(					_UpperCamelCase  ):





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__	:						int      		=  self.config_class(**self.inputs_dict      )
      self.parent.assertTrue(hasattr(A__					,			"hidden_sizes"      )      )
      self.parent.assertTrue(hasattr(A__					,			"num_attention_heads"      )      )







class   __lowerCAmelCase			:





   def __init__( self					,			snake_case					,			snake_case=13					,			snake_case=64					,			snake_case=3					,			snake_case=3					,			snake_case=2					,			snake_case=1					,			snake_case=16					,			snake_case=[128, 256, 384]					,			snake_case=[4, 6, 8]					,			snake_case=[2, 3, 4]					,			snake_case=[16, 16, 16]					,			snake_case=0					,			snake_case=[2, 2, 2]					,			snake_case=[2, 2, 2]					,			snake_case=0.02					,			snake_case=True					,			snake_case=True					,			snake_case=2					,			)				->  str:
      """simple docstring"""





      a__	:						Dict      		=  parent
      a__	:						str      		=  batch_size
      a__	:						Dict      		=  image_size
      a__	:						Optional[int]      		=  num_channels
      a__	:						List[str]      		=  kernel_size
      a__	:						List[str]      		=  stride
      a__	:						Union[str, Any]      		=  padding
      a__	:						int      		=  hidden_sizes
      a__	:						Union[str, Any]      		=  num_attention_heads
      a__	:						Union[str, Any]      		=  depths
      a__	:						List[str]      		=  key_dim
      a__	:						Any      		=  drop_path_rate
      a__	:						List[Any]      		=  patch_size
      a__	:						Optional[int]      		=  attention_ratio
      a__	:						Optional[Any]      		=  mlp_ratio
      a__	:						List[Any]      		=  initializer_range
      a__	:						List[Any]      		=  [
          ["""Subsample""", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
          ["""Subsample""", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
      ]
      a__	:						int      		=  is_training
      a__	:						List[str]      		=  use_labels
      a__	:						Tuple      		=  num_labels
      a__	:						int      		=  initializer_range





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      a__	:						str      		=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

      a__	:						Any      		=  None
      if self.use_labels:
         a__	:						Dict      		=  ids_tensor([self.batch_size]					,			self.num_labels      )

      a__	:						str      		=  self.get_config()
      return config, pixel_values, labels





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      return LevitConfig(
          image_size=self.image_size					,			num_channels=self.num_channels					,			kernel_size=self.kernel_size					,			stride=self.stride					,			padding=self.padding					,			patch_size=self.patch_size					,			hidden_sizes=self.hidden_sizes					,			num_attention_heads=self.num_attention_heads					,			depths=self.depths					,			key_dim=self.key_dim					,			drop_path_rate=self.drop_path_rate					,			mlp_ratio=self.mlp_ratio					,			attention_ratio=self.attention_ratio					,			initializer_range=self.initializer_range					,			down_ops=self.down_ops					,			)





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Optional[int]      		=  LevitModel(config=A__      )
      model.to(A__      )
      model.eval()
      a__	:						Tuple      		=  model(A__      )
      a__	:						str      		=  (self.image_size, self.image_size)
      a__	:						Any      		=  image_size[0], image_size[1]
      for _ in range(4      ):
         a__	:						List[Any]      		=  floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1      )
         a__	:						Union[str, Any]      		=  floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1      )
      self.parent.assertEqual(
          result.last_hidden_state.shape					,			(self.batch_size, ceil(height / 4      ) * ceil(width / 4      ), self.hidden_sizes[-1])					,			)





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      a__	:						Optional[Any]      		=  self.num_labels
      a__	:						Any      		=  LevitForImageClassification(A__      )
      model.to(A__      )
      model.eval()
      a__	:						List[Any]      		=  model(A__					,			labels=A__      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.num_labels)      )





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__	:						Tuple      		=  self.prepare_config_and_inputs()
      a__	:						List[Any]      		=  config_and_inputs
      a__	:						List[Any]      		=  {"""pixel_values""": pixel_values}
      return config, inputs_dict







@require_torch
class   __lowerCAmelCase			(					_UpperCamelCase		,_UpperCamelCase		,unittest.TestCase  ):
   _UpperCamelCase :						Union[str, Any]     		=				(
       (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
       if is_torch_available()
       else ()
   )
   _UpperCamelCase :						Union[str, Any]     		=				(
       {
           '''feature-extraction''': LevitModel,
           '''image-classification''': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
       }
       if is_torch_available()
       else {}
   )

   _UpperCamelCase :						List[str]     		=				False
   _UpperCamelCase :						List[Any]     		=				False
   _UpperCamelCase :						int     		=				False
   _UpperCamelCase :						Optional[int]     		=				False
   _UpperCamelCase :						Dict     		=				False





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Dict      		=  LevitModelTester(self      )
      a__	:						Optional[int]      		=  ConfigTester(self					,			config_class=A__					,			has_text_modality=A__					,			hidden_size=37      )





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      self.create_and_test_config_common_properties()
      self.config_tester.create_and_test_config_to_json_string()
      self.config_tester.create_and_test_config_to_json_file()
      self.config_tester.create_and_test_config_from_and_save_pretrained()
      self.config_tester.create_and_test_config_with_num_labels()
      self.config_tester.check_config_can_be_init_without_params()
      self.config_tester.check_config_arguments_init()





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      return





   @unittest.skip(reason="Levit does not use inputs_embeds"      )
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      pass





   @unittest.skip(reason="Levit does not support input and output embeddings"      )
   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      pass





   @unittest.skip(reason="Levit does not output attentions"      )
   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      pass





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      a__	:						Any      		=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         a__	:						str      		=  model_class(A__      )
         a__	:						List[Any]      		=  inspect.signature(model.forward      )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         a__	:						str      		=  [*signature.parameters.keys()]

         a__	:						Optional[Any]      		=  ["""pixel_values"""]
         self.assertListEqual(arg_names[:1]					,			A__      )





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      def check_hidden_states_output(snake_case					,			snake_case					,			snake_case      ):
         a__	:						str      		=  model_class(A__      )
         model.to(A__      )
         model.eval()

         with torch.no_grad():
            a__	:						List[Any]      		=  model(**self._prepare_for_class(A__					,			A__      )      )

         a__	:						List[Any]      		=  outputs.hidden_states

         a__	:						int      		=  len(self.model_tester.depths      ) + 1
         self.assertEqual(len(A__      )					,			A__      )

         a__	:						List[Any]      		=  (self.model_tester.image_size, self.model_tester.image_size)
         a__	:						List[str]      		=  image_size[0], image_size[1]
         for _ in range(4      ):
            a__	:						str      		=  floor(
                (
                    (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                    / self.model_tester.stride
                )
                + 1      )
            a__	:						List[str]      		=  floor(
                (
                    (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                    / self.model_tester.stride
                )
                + 1      )
         # verify the first hidden states (first block)
         self.assertListEqual(
             list(hidden_states[0].shape[-2:]      )					,			[
                 height * width,
                 self.model_tester.hidden_sizes[0],
             ]					,			)

      a__	:						str      		=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         a__	:						Union[str, Any]      		=  True
         check_hidden_states_output(A__					,			A__					,			A__      )

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         a__	:						Dict      		=  True

         check_hidden_states_output(A__					,			A__					,			A__      )





   @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."      )
   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      pass





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case=False      )				->  Tuple:
      """simple docstring"""





      a__	:						Tuple      		=  super()._prepare_for_class(A__					,			A__					,			return_labels=A__      )

      if return_labels:
         if model_class.__name__ == "LevitForImageClassificationWithTeacher":
            del inputs_dict["labels"]

      return inputs_dict





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__	:						List[str]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A__      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						int      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*A__      )





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      if not self.model_tester.is_training:
         return

      a__	:						Tuple      		=  self.model_tester.prepare_config_and_inputs_for_common()
      a__	:						Dict      		=  True

      for model_class in self.all_model_classes:
         # LevitForImageClassificationWithTeacher supports inference-only
         if (
             model_class in get_values(A__      )
             or model_class.__name__ == "LevitForImageClassificationWithTeacher"
         ):
            continue
         a__	:						Dict      		=  model_class(A__      )
         model.to(A__      )
         model.train()
         a__	:						Union[str, Any]      		=  self._prepare_for_class(A__					,			A__					,			return_labels=A__      )
         a__	:						List[Any]      		=  model(**A__      ).loss
         loss.backward()





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						Optional[int]      		=  self.model_tester.prepare_config_and_inputs_for_common()
      if not self.model_tester.is_training:
         return

      a__	:						Dict      		=  False
      a__	:						Any      		=  True

      for model_class in self.all_model_classes:
         if model_class in get_values(A__      ) or not model_class.supports_gradient_checkpointing:
            continue
         # LevitForImageClassificationWithTeacher supports inference-only
         if model_class.__name__ == "LevitForImageClassificationWithTeacher":
            continue
         a__	:						List[Any]      		=  model_class(A__      )
         model.gradient_checkpointing_enable()
         model.to(A__      )
         model.train()
         a__	:						Optional[Any]      		=  self._prepare_for_class(A__					,			A__					,			return_labels=A__      )
         a__	:						Any      		=  model(**A__      ).loss
         loss.backward()





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__	:						Optional[Any]      		=  self.model_tester.prepare_config_and_inputs_for_common()

      a__	:						Union[str, Any]      		=  [
          {"""title""": """multi_label_classification""", """num_labels""": 2, """dtype""": torch.float},
          {"""title""": """single_label_classification""", """num_labels""": 1, """dtype""": torch.long},
          {"""title""": """regression""", """num_labels""": 1, """dtype""": torch.float},
      ]

      for model_class in self.all_model_classes:
         if (
             model_class
             not in [
                 *get_values(A__      ),
             ]
             or model_class.__name__ == "LevitForImageClassificationWithTeacher"
         ):
            continue

         for problem_type in problem_types:
            with self.subTest(msg=F"""Testing {model_class} with {problem_type["title"]}"""      ):
               a__	:						Any      		=  problem_type["""title"""]
               a__	:						Any      		=  problem_type["""num_labels"""]

               a__	:						List[Any]      		=  model_class(A__      )
               model.to(A__      )
               model.train()

               a__	:						Optional[Any]      		=  self._prepare_for_class(A__					,			A__					,			return_labels=A__      )

               if problem_type["num_labels"] > 1:
                  a__	:						Any      		=  inputs["""labels"""].unsqueeze(1      ).repeat(1					,			problem_type["num_labels"]      )

               a__	:						Dict      		=  inputs["""labels"""].to(problem_type["dtype"]      )

               # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
               # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
               # they have the same size." which is a symptom something in wrong for the regression problem.
               # See https://github.com/huggingface/transformers/issues/11780
               with warnings.catch_warnings(record=A__      ) as warning_list:
                  a__	:						Dict      		=  model(**A__      ).loss
               for w in warning_list:
                  if "Using a target size that is different to the input size" in str(w.message      ):
                     raise ValueError(
                         F"""Something is going wrong in the regression problem: intercepted {w.message}"""      )

               loss.backward()





   @slow
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a__	:						Union[str, Any]      		=  LevitModel.from_pretrained(A__      )
         self.assertIsNotNone(A__      )







def       _A  (      ):
   a__	:						Union[str, Any]      		=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
   return image







@require_torch
@require_vision
class   __lowerCAmelCase			(					unittest.TestCase  ):





   @cached_property
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]      )





   @slow
   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Tuple      		=  LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]      ).to(
          A__      )

      a__	:						Optional[int]      		=  self.default_image_processor
      a__	:						List[Any]      		=  prepare_img()
      a__	:						Dict      		=  image_processor(images=A__					,			return_tensors="pt"      ).to(A__      )

      # forward pass
      with torch.no_grad():
         a__	:						Optional[int]      		=  model(**A__      )

      # verify the logits
      a__	:						List[Any]      		=  torch.Size((1, 1_000)      )
      self.assertEqual(outputs.logits.shape					,			A__      )

      a__	:						Any      		=  torch.tensor([1.0_448, -0.3_745, -1.8_317]      ).to(A__      )

      self.assertTrue(torch.allclose(outputs.logits[0, :3]					,			A__					,			atol=1E-4      )      )








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  (      lowerCamelCase   ):
   a__	:						List[str]      		=  []
   if isinstance(lowerCamelCase							, lowerCamelCase   ):
      for v in tree.values():
         shapes.extend(_fetch_dims(lowerCamelCase   )   )
   elif isinstance(lowerCamelCase							, (list, tuple)   ):
      for t in tree:
         shapes.extend(_fetch_dims(lowerCamelCase   )   )
   elif isinstance(lowerCamelCase							, torch.Tensor   ):
      shapes.append(tree.shape   )
   else:
      raise ValueError("Not supported"   )

   return shapes






@torch.jit.ignore
def       _A  (      lowerCamelCase							, lowerCamelCase   ):
   a__	:						List[str]      		=  []
   for d in reversed(lowerCamelCase   ):
      idx.append(flat_idx % d   )
      a__	:						Union[str, Any]      		=  flat_idx // d

   return tuple(reversed(lowerCamelCase   )   )






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

   # start_edges and end_edges both indicate whether, starting from any given
   # dimension, the start/end index is at the top/bottom edge of the
   # corresponding tensor, modeled as a tree
   def reduce_edge_list(lowerCamelCase   ) -> None:
      a__	:						int      		=  True
      for i in range(len(lowerCamelCase   )   ):
         a__	:						Optional[Any]      		=  -1 * (i + 1)
         l[reversed_idx] &= tally
         a__	:						Tuple      		=  l[reversed_idx]

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

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

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

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

   a__	:						Tuple[slice, ...]      		=  tuple(lowerCamelCase   )
   a__	:						Optional[Any]      		=  len(lowerCamelCase   )

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

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

      a__	:						Optional[Any]      		=  start[divergence_idx]
      return tuple(
          path + (slice(lowerCamelCase							, 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

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

   # Get an ordered list of slices to perform
   a__	:						str      		=  _get_minimal_slice_set(
       lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, )

   a__	:						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  (      lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, lowerCamelCase = False							, lowerCamelCase = None							, lowerCamelCase = False							, ):
   if not (len(lowerCamelCase   ) > 0):
      raise ValueError("Must provide at least one input"   )

   a__	:						str      		=  [shape[:no_batch_dims] for shape in _fetch_dims(lowerCamelCase   )]
   a__	:						Dict      		=  tuple([max(lowerCamelCase   ) for s in zip(*lowerCamelCase   )]   )

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

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

   a__	:						Optional[Any]      		=  1
   for d in orig_batch_dims:
      flat_batch_dim *= d

   a__	:						Tuple      		=  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

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

      a__	:						Dict[str, Any]      		=  tensor_tree_map(lowerCamelCase							, lowerCamelCase   )

      # Run the layer on the chunk
      a__	:						Any      		=  layer(**lowerCamelCase   )

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

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

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

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

      i += chunk_size

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

   return out







class   __lowerCAmelCase			:





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





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





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  int:
      """simple docstring"""





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

      if min_chunk_size >= self.max_chunk_size:
         return min_chunk_size

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

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

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

      return candidates[min_viable_chunk_size_index]





   def 	_snake_case ( self					,			snake_case					,			snake_case      )				->  bool:
      """simple docstring"""





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

      return consistent





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			)				->  int:
      """simple docstring"""





      a__	:						List[Any]      		=  True
      a__	:						tuple      		=  tree_map(lambda snake_case      : a.shape if isinstance(snake_case					,			torch.Tensor      ) else a					,			snake_case					,			snake_case      )
      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(snake_case      )
         a__	:						Union[str, Any]      		=  self._compare_arg_caches(self.cached_arg_data					,			snake_case      )
      else:
         # Otherwise, we can reuse the precomputed value
         a__	:						Optional[int]      		=  False

      if not consistent:
         a__	:						List[str]      		=  self._determine_favorable_chunk_size(
             snake_case					,			snake_case					,			snake_case					,			)
         a__	:						List[str]      		=  arg_data

      assert self.cached_chunk_size is not None

      return self.cached_chunk_size

