






import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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 MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
       from transformers.models.mobilevitva.modeling_mobilevitva import (
           MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
           make_divisible,
       )


if is_vision_available():
       from PIL import Image

       from transformers import MobileViTImageProcessor





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  self.config_class(**self.inputs_dict )
           self.parent.assertTrue(hasattr(_SCREAMING_SNAKE_CASE   ,						'width_multiplier' ) )





class 					A__     :

     '''simple docstring'''






     def __init__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=13   ,						_SCREAMING_SNAKE_CASE      :    Dict=64   ,						_SCREAMING_SNAKE_CASE      :    List[str]=2   ,						_SCREAMING_SNAKE_CASE      :    int=3   ,						_SCREAMING_SNAKE_CASE      :    Any="swish"   ,						_SCREAMING_SNAKE_CASE      :    List[str]=3   ,						_SCREAMING_SNAKE_CASE      :    Dict=32   ,						_SCREAMING_SNAKE_CASE      :    Any=0.1   ,						_SCREAMING_SNAKE_CASE      :    Dict=0.0_2   ,						_SCREAMING_SNAKE_CASE      :    List[str]=True   ,						_SCREAMING_SNAKE_CASE      :    Dict=True   ,						_SCREAMING_SNAKE_CASE      :    int=10   ,						_SCREAMING_SNAKE_CASE      :    Dict=None   ,						_SCREAMING_SNAKE_CASE      :    List[str]=0.2_5   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=0.0   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=0.0   ,						):






           """simple docstring"""




           UpperCamelCase					       =  parent
           UpperCamelCase					       =  batch_size
           UpperCamelCase					       =  image_size
           UpperCamelCase					       =  patch_size
           UpperCamelCase					       =  num_channels
           UpperCamelCase					       =  make_divisible(512 * width_multiplier   ,						divisor=8 )
           UpperCamelCase					       =  hidden_act
           UpperCamelCase					       =  conv_kernel_size
           UpperCamelCase					       =  output_stride
           UpperCamelCase					       =  classifier_dropout_prob
           UpperCamelCase					       =  use_labels
           UpperCamelCase					       =  is_training
           UpperCamelCase					       =  num_labels
           UpperCamelCase					       =  initializer_range
           UpperCamelCase					       =  scope
           UpperCamelCase					       =  width_multiplier
           UpperCamelCase					       =  ffn_dropout
           UpperCamelCase					       =  attn_dropout







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

           UpperCamelCase					       =  None
           UpperCamelCase					       =  None
           if self.use_labels:
                 UpperCamelCase					       =  ids_tensor([self.batch_size]   ,						self.num_labels )
                 UpperCamelCase					       =  ids_tensor([self.batch_size, self.image_size, self.image_size]   ,						self.num_labels )

           UpperCamelCase					       =  self.get_config()

           return config, pixel_values, labels, pixel_labels







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           return MobileViTVaConfig(
               image_size=self.image_size   ,						patch_size=self.patch_size   ,						num_channels=self.num_channels   ,						hidden_act=self.hidden_act   ,						conv_kernel_size=self.conv_kernel_size   ,						output_stride=self.output_stride   ,						classifier_dropout_prob=self.classifier_dropout_prob   ,						initializer_range=self.initializer_range   ,						width_multiplier=self.width_multiplier   ,						ffn_dropout=self.ffn_dropout_prob   ,						attn_dropout=self.attn_dropout_prob   ,						)







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  MobileViTVaModel(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(
               result.last_hidden_state.shape   ,						(
                   self.batch_size,
                   self.last_hidden_size,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )   ,						)







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  self.num_labels
           UpperCamelCase					       =  MobileViTVaForImageClassification(_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_labels) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  self.num_labels
           UpperCamelCase					       =  MobileViTVaForSemanticSegmentation(_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(
               result.logits.shape   ,						(
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )   ,						)
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(
               result.logits.shape   ,						(
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )   ,						)







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           UpperCamelCase					       =  self.prepare_config_and_inputs()
           UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  config_and_inputs
           UpperCamelCase					       =  {'pixel_values': pixel_values}
           return config, inputs_dict





@require_torch
class 					A__     (     __snake_case   ,    __snake_case   ,    unittest.TestCase ):

     '''simple docstring'''






     snake_case__					=					(
         (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
         if is_torch_available()
         else ()
     )

     snake_case__					=					(
         {
             """feature-extraction""": MobileViTVaModel,
             """image-classification""": MobileViTVaForImageClassification,
             """image-segmentation""": MobileViTVaForSemanticSegmentation,
         }
         if is_torch_available()
         else {}
     )

     snake_case__					=					False
     snake_case__					=					False
     snake_case__					=					False
     snake_case__					=					False







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  MobileViTVaModelTester(self )
           UpperCamelCase					       =  MobileViTVaConfigTester(self   ,						config_class=_SCREAMING_SNAKE_CASE   ,						has_text_modality=_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           self.config_tester.run_common_tests()







     @unittest.skip(reason='MobileViTV2 does not use inputs_embeds' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int] ):






           """simple docstring"""




           pass







     @unittest.skip(reason='MobileViTV2 does not support input and output embeddings' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           pass







     @unittest.skip(reason='MobileViTV2 does not output attentions' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           pass







     @require_torch_multi_gpu
     @unittest.skip(reason='Got `CUDA error: misaligned address` for tests after this one being run.' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           pass







     @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           pass







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           UpperCamelCase       ,  UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs_for_common()

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

                 UpperCamelCase					       =  ['pixel_values']
                 self.assertListEqual(arg_names[:1]   ,						_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           def check_hidden_states_output(_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):
                 UpperCamelCase					       =  model_class(_SCREAMING_SNAKE_CASE )
                 model.to(_SCREAMING_SNAKE_CASE )
                 model.eval()

                 with torch.no_grad():
                       UpperCamelCase					       =  model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) )

                 UpperCamelCase					       =  outputs.hidden_states

                 UpperCamelCase					       =  5
                 self.assertEqual(len(_SCREAMING_SNAKE_CASE )   ,						_SCREAMING_SNAKE_CASE )

                 # MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
                 # with the width and height being successively divided by 2.
                 UpperCamelCase					       =  2
                 for i in range(len(_SCREAMING_SNAKE_CASE ) ):
                       self.assertListEqual(
                           list(hidden_states[i].shape[-2:] )   ,						[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]   ,						)
                       divisor *= 2

                 self.assertEqual(self.model_tester.output_stride   ,						divisor // 2 )

           UpperCamelCase       ,  UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 UpperCamelCase					       =  True
                 check_hidden_states_output(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 UpperCamelCase					       =  True

                 check_hidden_states_output(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int] ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*_SCREAMING_SNAKE_CASE )







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 UpperCamelCase					       =  MobileViTVaModel.from_pretrained(_SCREAMING_SNAKE_CASE )
                 self.assertIsNotNone(_SCREAMING_SNAKE_CASE )

def 		lowercase__  (				) ->			int:
      """simple docstring"""

      UpperCamelCase					       =  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
      return image





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

     '''simple docstring'''






     @cached_property
     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           return (
               MobileViTImageProcessor.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256' )
               if is_vision_available()
               else None
           )







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  MobileViTVaForImageClassification.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256' ).to(
               _SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  self.default_image_processor
           UpperCamelCase					       =  prepare_img()
           UpperCamelCase					       =  image_processor(images=_SCREAMING_SNAKE_CASE   ,						return_tensors='pt' ).to(_SCREAMING_SNAKE_CASE )

           # forward pass
           with torch.no_grad():
                 UpperCamelCase					       =  model(**_SCREAMING_SNAKE_CASE )

           # verify the logits
           UpperCamelCase					       =  torch.Size((1, 1000) )
           self.assertEqual(outputs.logits.shape   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  torch.tensor([-1.6_3_3_6E0_0, -7.3_2_0_4E-0_2, -5.1_8_8_3E-0_1] ).to(_SCREAMING_SNAKE_CASE )

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







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )
           UpperCamelCase					       =  model.to(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )

           UpperCamelCase					       =  prepare_img()
           UpperCamelCase					       =  image_processor(images=_SCREAMING_SNAKE_CASE   ,						return_tensors='pt' ).to(_SCREAMING_SNAKE_CASE )

           # forward pass
           with torch.no_grad():
                 UpperCamelCase					       =  model(**_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  outputs.logits

           # verify the logits
           UpperCamelCase					       =  torch.Size((1, 21, 32, 32) )
           self.assertEqual(logits.shape   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  torch.tensor(
               [
                   [[7.0_8_6_3, 7.1_5_2_5, 6.8_2_0_1], [6.6_9_3_1, 6.8_7_7_0, 6.8_9_3_3], [6.2_9_7_8, 7.0_3_6_6, 6.9_6_3_6]],
                   [[-3.7_1_3_4, -3.6_7_1_2, -3.6_6_7_5], [-3.5_8_2_5, -3.3_5_4_9, -3.4_7_7_7], [-3.3_4_3_5, -3.3_9_7_9, -3.2_8_5_7]],
                   [[-2.9_3_2_9, -2.8_0_0_3, -2.7_3_6_9], [-3.0_5_6_4, -2.4_7_8_0, -2.0_2_0_7], [-2.6_8_8_9, -1.9_2_9_8, -1.7_6_4_0]],
               ]   ,						device=_SCREAMING_SNAKE_CASE   ,						)

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







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int] ):






           """simple docstring"""




           UpperCamelCase					       =  MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )
           UpperCamelCase					       =  model.to(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )

           UpperCamelCase					       =  prepare_img()
           UpperCamelCase					       =  image_processor(images=_SCREAMING_SNAKE_CASE   ,						return_tensors='pt' ).to(_SCREAMING_SNAKE_CASE )

           # forward pass
           with torch.no_grad():
                 UpperCamelCase					       =  model(**_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  outputs.logits.detach().cpu()

           UpperCamelCase					       =  image_processor.post_process_semantic_segmentation(outputs=_SCREAMING_SNAKE_CASE   ,						target_sizes=[(50, 60)] )
           UpperCamelCase					       =  torch.Size((50, 60) )
           self.assertEqual(segmentation[0].shape   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  image_processor.post_process_semantic_segmentation(outputs=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  torch.Size((32, 32) )
           self.assertEqual(segmentation[0].shape   ,						_SCREAMING_SNAKE_CASE )











from manim import *





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  Rectangle(height=0.5   ,						width=0.5 )
           UpperCamelCase					       =  Rectangle(height=0.4_6   ,						width=0.4_6 ).set_stroke(width=0 )

           UpperCamelCase					       =  Rectangle(height=0.2_5   ,						width=0.2_5 )

           UpperCamelCase					       =  [mem.copy() for i in range(6 )]
           UpperCamelCase					       =  [mem.copy() for i in range(6 )]
           UpperCamelCase					       =  VGroup(*_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  VGroup(*_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  VGroup(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  Text('CPU'   ,						font_size=24 )
           UpperCamelCase					       =  Group(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0.5   ,						aligned_edge=_SCREAMING_SNAKE_CASE )
           cpu.move_to([-2.5, -0.5, 0] )
           self.add(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  [mem.copy() for i in range(4 )]
           UpperCamelCase					       =  VGroup(*_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  Text('GPU'   ,						font_size=24 )
           UpperCamelCase					       =  Group(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0.5   ,						aligned_edge=_SCREAMING_SNAKE_CASE )
           gpu.move_to([-1, -1, 0] )
           self.add(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  [mem.copy() for i in range(6 )]
           UpperCamelCase					       =  VGroup(*_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )

           UpperCamelCase					       =  Text('Model'   ,						font_size=24 )
           UpperCamelCase					       =  Group(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0.5   ,						aligned_edge=_SCREAMING_SNAKE_CASE )
           model.move_to([3, -1.0, 0] )
           self.add(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  []
           UpperCamelCase					       =  []

           for i, rect in enumerate(_SCREAMING_SNAKE_CASE ):
                 UpperCamelCase					       =  fill.copy().set_fill(_SCREAMING_SNAKE_CASE   ,						opacity=0.8 )
                 target.move_to(_SCREAMING_SNAKE_CASE )
                 model_arr.append(_SCREAMING_SNAKE_CASE )

                 UpperCamelCase					       =  Rectangle(height=0.4_6   ,						width=0.4_6 ).set_stroke(width=0.0 ).set_fill(_SCREAMING_SNAKE_CASE   ,						opacity=0.8 )
                 cpu_target.move_to(cpu_left_col_base[i] )
                 model_cpu_arr.append(_SCREAMING_SNAKE_CASE )

           self.add(*_SCREAMING_SNAKE_CASE   ,						*_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  [meta_mem.copy() for i in range(6 )]
           UpperCamelCase					       =  [meta_mem.copy() for i in range(6 )]
           UpperCamelCase					       =  VGroup(*_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  VGroup(*_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  VGroup(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0 )
           UpperCamelCase					       =  Text('Disk'   ,						font_size=24 )
           UpperCamelCase					       =  Group(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ).arrange(_SCREAMING_SNAKE_CASE   ,						buff=0.5   ,						aligned_edge=_SCREAMING_SNAKE_CASE )
           disk.move_to([-4, -1.2_5, 0] )
           self.add(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  Square(side_length=2.2 )
           key.move_to([-5, 2, 0] )

           UpperCamelCase					       =  MarkupText(
               f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'   ,						font_size=18   ,						)

           key_text.move_to([-5, 2.4, 0] )

           self.add(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  MarkupText(
               f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'   ,						font_size=18   ,						)

           blue_text.next_to(_SCREAMING_SNAKE_CASE   ,						DOWN * 2.4   ,						aligned_edge=key_text.get_left() )
           self.add(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  MarkupText(
               f'Now watch as an input is passed through the model\nand how the memory is utilized and handled.'   ,						font_size=24   ,						)
           step_a.move_to([2, 2, 0] )

           self.play(Write(_SCREAMING_SNAKE_CASE ) )

           UpperCamelCase					       =  Square(0.3 )
           input.set_fill(_SCREAMING_SNAKE_CASE   ,						opacity=1.0 )
           input.set_stroke(width=0.0 )
           input.next_to(model_base[0]   ,						_SCREAMING_SNAKE_CASE   ,						buff=0.5 )

           self.play(Write(_SCREAMING_SNAKE_CASE ) )

           input.generate_target()
           input.target.next_to(model_arr[0]   ,						direction=_SCREAMING_SNAKE_CASE   ,						buff=0.0_2 )
           self.play(MoveToTarget(_SCREAMING_SNAKE_CASE ) )

           self.play(FadeOut(_SCREAMING_SNAKE_CASE ) )

           UpperCamelCase					       =  Arrow(start=_SCREAMING_SNAKE_CASE   ,						end=_SCREAMING_SNAKE_CASE   ,						color=_SCREAMING_SNAKE_CASE   ,						buff=0.5 )
           a.next_to(model_arr[0].get_left()   ,						_SCREAMING_SNAKE_CASE   ,						buff=0.2 )

           model_cpu_arr[0].generate_target()
           model_cpu_arr[0].target.move_to(gpu_rect[0] )

           UpperCamelCase					       =  MarkupText(
               f'As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back.'   ,						font_size=24   ,						)
           step_a.move_to([2, 2, 0] )

           self.play(Write(_SCREAMING_SNAKE_CASE   ,						run_time=3 ) )

           UpperCamelCase					       =  {'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.0_2}

           self.play(
               Write(_SCREAMING_SNAKE_CASE )   ,						Circumscribe(model_arr[0]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(model_cpu_arr[0]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(gpu_rect[0]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						)
           self.play(MoveToTarget(model_cpu_arr[0] ) )

           UpperCamelCase					       =  a.copy()
           for i in range(6 ):
                 a_c.next_to(model_arr[i].get_right() + 0.0_2   ,						_SCREAMING_SNAKE_CASE   ,						buff=0.2 )

                 input.generate_target()
                 input.target.move_to(model_arr[i].get_right() + 0.0_2 )

                 UpperCamelCase					       =  AnimationGroup(
                     FadeOut(_SCREAMING_SNAKE_CASE   ,						run_time=0.5 )   ,						MoveToTarget(_SCREAMING_SNAKE_CASE   ,						run_time=0.5 )   ,						FadeIn(_SCREAMING_SNAKE_CASE   ,						run_time=0.5 )   ,						lag_ratio=0.2 )

                 self.play(_SCREAMING_SNAKE_CASE )

                 model_cpu_arr[i].generate_target()
                 model_cpu_arr[i].target.move_to(cpu_left_col_base[i] )

                 if i < 5:
                       model_cpu_arr[i + 1].generate_target()
                       model_cpu_arr[i + 1].target.move_to(gpu_rect[0] )
                       if i >= 1:
                             UpperCamelCase					       =  0.7

                       self.play(
                           Circumscribe(model_arr[i]   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(cpu_left_col_base[i]   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(cpu_left_col_base[i + 1]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(gpu_rect[0]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(model_arr[i + 1]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						)
                       if i < 1:
                             self.play(
                                 MoveToTarget(model_cpu_arr[i] )   ,						MoveToTarget(model_cpu_arr[i + 1] )   ,						)
                       else:
                             self.play(
                                 MoveToTarget(model_cpu_arr[i]   ,						run_time=0.7 )   ,						MoveToTarget(model_cpu_arr[i + 1]   ,						run_time=0.7 )   ,						)
                 else:
                       model_cpu_arr[i].generate_target()
                       model_cpu_arr[i].target.move_to(cpu_left_col_base[-1] )
                       input.generate_target()
                       input.target.next_to(model_arr[-1].get_right()   ,						RIGHT + 0.0_2   ,						buff=0.2 )

                       self.play(
                           Circumscribe(model_arr[-1]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(cpu_left_col_base[-1]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						Circumscribe(gpu_rect[0]   ,						color=_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )   ,						)

                       self.play(MoveToTarget(model_cpu_arr[i] ) )

                 UpperCamelCase					       =  a_c
                 UpperCamelCase					       =  a_c.copy()

           input.generate_target()
           input.target.next_to(model_base[-1]   ,						RIGHT + 0.0_2   ,						buff=0.5 )
           self.play(
               FadeOut(_SCREAMING_SNAKE_CASE )   ,						FadeOut(_SCREAMING_SNAKE_CASE   ,						run_time=0.5 )   ,						)

           UpperCamelCase					       =  MarkupText(f'Inference on a model too large for GPU memory\nis successfully completed.'   ,						font_size=24 )
           step_a.move_to([2, 2, 0] )

           self.play(Write(_SCREAMING_SNAKE_CASE   ,						run_time=3 )   ,						MoveToTarget(_SCREAMING_SNAKE_CASE ) )

           self.wait()



