


import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
   import os

   import jax.numpy as jnp
   from jax import jit

   from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
   from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

   _SCREAMING_SNAKE_CASE			:       Tuple      =						"""0.12"""  # assumed parallelism: 8


if is_torch_available():
   import torch







def     _lowercase							(  __lowerCamelCase      :       Optional[Any]		,__lowerCamelCase      :       Union[str, Any]		,__lowerCamelCase      :       Optional[int]=None      )	-> int:





     '''simple docstring'''





     if rng is None:
          UpperCamelCase__				:	Dict										=  random.Random()

     UpperCamelCase__				:	Optional[Any]										=  1
     for dim in shape:
          total_dims *= dim

     UpperCamelCase__				:	List[str]										=  []
     for _ in range(__lowerCAmelCase      ):
          values.append(rng.randint(0		,vocab_size - 1      )      )

     UpperCamelCase__				:	int										=  np.array(__lowerCAmelCase		,dtype=jnp.intaa      ).reshape(__lowerCAmelCase      )

     return output







def     _lowercase							(  __lowerCamelCase      :       Dict		,__lowerCamelCase      :       List[Any]=None      )	-> List[str]:





     '''simple docstring'''





     UpperCamelCase__				:	Union[str, Any]										=  ids_tensor(__lowerCAmelCase		,vocab_size=2		,rng=__lowerCAmelCase      )
     # make sure that at least one token is attended to for each batch
     UpperCamelCase__				:	int										=  1
     return attn_mask

@require_flax
class 			UpperCamelCase__			:
      a__ :  int		    =    None
      a__ :  Dict		    =    ()






      def 						__lowercase(  self       : Dict )	->			str:
           UpperCamelCase__				:	Any										=  self.model_tester.prepare_config_and_inputs_for_common()

           # cut to half length & take max batch_size 3
           UpperCamelCase__				:	List[str]										=  2
           UpperCamelCase__				:	List[str]										=  inputs["""input_ids"""].shape[-1] // 2
           UpperCamelCase__				:	int										=  inputs["""input_ids"""][:max_batch_size, :sequence_length]

           UpperCamelCase__				:	Union[str, Any]										=  jnp.ones_like(lowerCamelCase__ )
           UpperCamelCase__				:	Any										=  attention_mask[:max_batch_size, :sequence_length]

           # generate max 5 tokens
           UpperCamelCase__				:	Optional[int]										=  input_ids.shape[-1] + 5
           if config.eos_token_id is not None and config.pad_token_id is None:
                # hack to allow generate for models such as GPT2 as is done in `generate()`
                UpperCamelCase__				:	str										=  config.eos_token_id
           return config, input_ids, attention_mask, max_length






      @is_pt_flax_cross_test
      def 						__lowercase(  self       : int )	->			List[str]:
           UpperCamelCase__				:	Optional[int]										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	List[str]										=  False
           UpperCamelCase__				:	Dict										=  max_length
           UpperCamelCase__				:	Optional[int]										=  0

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Any										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	List[Any]										=  model_class.__name__[4:]  # Skip the "Flax" at the beginning
                UpperCamelCase__				:	List[str]										=  getattr(lowerCamelCase__,		lowerCamelCase__ )
                UpperCamelCase__				:	Any										=  pt_model_class(lowerCamelCase__ ).eval()
                UpperCamelCase__				:	Optional[Any]										=  load_flax_weights_in_pytorch_model(lowerCamelCase__,		flax_model.params )

                UpperCamelCase__				:	Any										=  flax_model.generate(lowerCamelCase__ ).sequences
                UpperCamelCase__				:	str										=  pt_model.generate(torch.tensor(lowerCamelCase__,		dtype=torch.long ) )

                if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
                     UpperCamelCase__				:	Optional[int]										=  flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

                self.assertListEqual(pt_generation_outputs.numpy().tolist(),		flax_generation_outputs.tolist() )






      def 						__lowercase(  self       : str )	->			Optional[Any]:
           UpperCamelCase__				:	Optional[Any]										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	Optional[Any]										=  False
           UpperCamelCase__				:	int										=  max_length

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Union[str, Any]										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	Tuple										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	List[str]										=  jit(model.generate )
                UpperCamelCase__				:	Optional[Any]										=  jit_generate(lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : Tuple )	->			List[str]:
           UpperCamelCase__				:	str										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	Optional[Any]										=  True
           UpperCamelCase__				:	List[str]										=  max_length

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Union[str, Any]										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	List[str]										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	Union[str, Any]										=  jit(model.generate )
                UpperCamelCase__				:	Any										=  jit_generate(lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : Optional[Any] )	->			Any:
           UpperCamelCase__				:	Optional[Any]										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	Any										=  False
           UpperCamelCase__				:	Optional[Any]										=  max_length
           UpperCamelCase__				:	Tuple										=  2

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Any										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	Dict										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	str										=  jit(model.generate )
                UpperCamelCase__				:	int										=  jit_generate(lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : str )	->			Any:
           UpperCamelCase__				:	str										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	List[str]										=  False
           UpperCamelCase__				:	Union[str, Any]										=  max_length
           UpperCamelCase__				:	Any										=  2
           UpperCamelCase__				:	Tuple										=  2

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	List[Any]										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	Optional[int]										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[0],		input_ids.shape[0] * config.num_return_sequences )






      def 						__lowercase(  self       : str )	->			Optional[Any]:
           UpperCamelCase__				:	List[str]										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	Dict										=  True
           UpperCamelCase__				:	Dict										=  max_length
           UpperCamelCase__				:	Optional[Any]										=  0.8
           UpperCamelCase__				:	int										=  10
           UpperCamelCase__				:	List[Any]										=  0.3
           UpperCamelCase__				:	str										=  1
           UpperCamelCase__				:	Union[str, Any]										=  8
           UpperCamelCase__				:	int										=  9

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Optional[int]										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	Optional[Any]										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	Dict										=  jit(model.generate )
                UpperCamelCase__				:	Any										=  jit_generate(lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : Tuple )	->			str:
           UpperCamelCase__				:	Optional[int]										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	Dict										=  max_length
           UpperCamelCase__				:	Any										=  1
           UpperCamelCase__				:	List[str]										=  8
           UpperCamelCase__				:	Union[str, Any]										=  9

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	str										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	List[str]										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	str										=  jit(model.generate )
                UpperCamelCase__				:	List[str]										=  jit_generate(lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : List[Any] )	->			Dict:
           UpperCamelCase__				:	List[Any]										=  self._get_input_ids_and_config()
           UpperCamelCase__				:	Dict										=  max_length
           UpperCamelCase__				:	Dict										=  2
           UpperCamelCase__				:	Dict										=  1
           UpperCamelCase__				:	List[str]										=  8
           UpperCamelCase__				:	List[Any]										=  9

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Tuple										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	int										=  model.generate(lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	Tuple										=  jit(model.generate )
                UpperCamelCase__				:	List[Any]										=  jit_generate(lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : List[str] )	->			Optional[int]:
           UpperCamelCase__				:	Optional[Any]										=  self._get_input_ids_and_config()

           # pad attention mask on the left
           UpperCamelCase__				:	Dict										=  attention_mask.at[(0, 0)].set(0 )

           UpperCamelCase__				:	List[str]										=  False
           UpperCamelCase__				:	str										=  max_length

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Optional[int]										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	Union[str, Any]										=  model.generate(lowerCamelCase__,		attention_mask=lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	str										=  jit(model.generate )
                UpperCamelCase__				:	int										=  jit_generate(lowerCamelCase__,		attention_mask=lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : int )	->			int:
           UpperCamelCase__				:	List[str]										=  self._get_input_ids_and_config()

           # pad attention mask on the left
           UpperCamelCase__				:	int										=  attention_mask.at[(0, 0)].set(0 )

           UpperCamelCase__				:	Union[str, Any]										=  True
           UpperCamelCase__				:	List[str]										=  max_length

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	str										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	Union[str, Any]										=  model.generate(lowerCamelCase__,		attention_mask=lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	Dict										=  jit(model.generate )
                UpperCamelCase__				:	List[str]										=  jit_generate(lowerCamelCase__,		attention_mask=lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )






      def 						__lowercase(  self       : List[Any] )	->			Optional[int]:
           UpperCamelCase__				:	str										=  self._get_input_ids_and_config()

           # pad attention mask on the left
           UpperCamelCase__				:	Union[str, Any]										=  attention_mask.at[(0, 0)].set(0 )

           UpperCamelCase__				:	int										=  2
           UpperCamelCase__				:	List[Any]										=  max_length

           for model_class in self.all_generative_model_classes:
                UpperCamelCase__				:	Optional[int]										=  model_class(lowerCamelCase__ )

                UpperCamelCase__				:	int										=  model.generate(lowerCamelCase__,		attention_mask=lowerCamelCase__ ).sequences
                self.assertEqual(generation_outputs.shape[-1],		lowerCamelCase__ )

                UpperCamelCase__				:	Optional[Any]										=  jit(model.generate )
                UpperCamelCase__				:	int										=  jit_generate(lowerCamelCase__,		attention_mask=lowerCamelCase__ ).sequences

                self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist() )

@require_flax
class 			UpperCamelCase__			(			unittest.TestCase      ):






      def 						__lowercase(  self       : List[Any] )	->			Any:
           UpperCamelCase__				:	str										=  AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert''' )
           UpperCamelCase__				:	Union[str, Any]										=  FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''' )

           UpperCamelCase__				:	Optional[int]										=  """Hello world"""
           UpperCamelCase__				:	Dict										=  tokenizer(lowerCamelCase__,		return_tensors='''np''' ).input_ids

           # typos are quickly detected (the correct argument is `do_sample`)
           with self.assertRaisesRegex(lowerCamelCase__,		'''do_samples''' ):
                model.generate(lowerCamelCase__,		do_samples=lowerCamelCase__ )

           # arbitrary arguments that will not be used anywhere are also not accepted
           with self.assertRaisesRegex(lowerCamelCase__,		'''foo''' ):
                UpperCamelCase__				:	int										=  {"""foo""": """bar"""}
                model.generate(lowerCamelCase__,		**lowerCamelCase__ )












from manim import *







class 		_A						( UpperCAmelCase_ ):

  def    a					(    self						:	Dict			):
     """simple docstring"""

     __UpperCamelCase	:			Tuple					=					Rectangle(height=0.5   ,				width=0.5			)
     __UpperCamelCase	:			Tuple					=					Rectangle(height=0.46   ,				width=0.46			).set_stroke(width=0			)

     __UpperCamelCase	:			Optional[Any]					=					Rectangle(height=0.25   ,				width=0.25			)

     __UpperCamelCase	:			Any					=					[mem.copy() for i in range(6			)]
     __UpperCamelCase	:			str					=					[mem.copy() for i in range(6			)]
     __UpperCamelCase	:			int					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Optional[int]					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Dict					=					VGroup(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			List[Any]					=					Text("""CPU"""   ,				font_size=24			)
     __UpperCamelCase	:			Any					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     cpu.move_to([-2.5, -0.5, 0]			)
     self.add(lowerCamelCase__			)

     __UpperCamelCase	:			Tuple					=					[mem.copy() for i in range(4			)]
     __UpperCamelCase	:			Optional[Any]					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Optional[Any]					=					Text("""GPU"""   ,				font_size=24			)
     __UpperCamelCase	:			Union[str, Any]					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     gpu.move_to([-1, -1, 0]			)
     self.add(lowerCamelCase__			)

     __UpperCamelCase	:			Dict					=					[mem.copy() for i in range(6			)]
     __UpperCamelCase	:			Tuple					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)

     __UpperCamelCase	:			Tuple					=					Text("""Model"""   ,				font_size=24			)
     __UpperCamelCase	:			List[str]					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     model.move_to([3, -1.0, 0]			)
     self.add(lowerCamelCase__			)

     __UpperCamelCase	:			Tuple					=					[]
     __UpperCamelCase	:			str					=					[]

     for i, rect in enumerate(lowerCamelCase__			):
        __UpperCamelCase	:			Optional[Any]					=					fill.copy().set_fill(lowerCamelCase__   ,				opacity=0.8			)
        target.move_to(lowerCamelCase__			)
        model_arr.append(lowerCamelCase__			)

        __UpperCamelCase	:			Tuple					=					Rectangle(height=0.46   ,				width=0.46			).set_stroke(width=0.0			).set_fill(lowerCamelCase__   ,				opacity=0.8			)
        cpu_target.move_to(cpu_left_col_base[i]			)
        model_cpu_arr.append(lowerCamelCase__			)

     self.add(*lowerCamelCase__   ,				*lowerCamelCase__			)

     __UpperCamelCase	:			Union[str, Any]					=					[meta_mem.copy() for i in range(6			)]
     __UpperCamelCase	:			Any					=					[meta_mem.copy() for i in range(6			)]
     __UpperCamelCase	:			int					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Union[str, Any]					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			str					=					VGroup(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Optional[Any]					=					Text("""Disk"""   ,				font_size=24			)
     __UpperCamelCase	:			Optional[Any]					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     disk.move_to([-4, -1.25, 0]			)
     self.add(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			Union[str, Any]					=					Square(side_length=2.2			)
     key.move_to([-5, 2, 0]			)

     __UpperCamelCase	:			Optional[Any]					=					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(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			Optional[int]					=					MarkupText(
         f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'   ,				font_size=18   ,				)

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

     __UpperCamelCase	:			Any					=					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(lowerCamelCase__			)			)

     __UpperCamelCase	:			List[str]					=					Square(0.3			)
     input.set_fill(lowerCamelCase__   ,				opacity=1.0			)
     input.set_stroke(width=0.0			)
     input.next_to(model_base[0]   ,				lowerCamelCase__   ,				buff=0.5			)

     self.play(Write(lowerCamelCase__			)			)

     input.generate_target()
     input.target.next_to(model_arr[0]   ,				direction=lowerCamelCase__   ,				buff=0.02			)
     self.play(MoveToTarget(lowerCamelCase__			)			)

     self.play(FadeOut(lowerCamelCase__			)			)

     __UpperCamelCase	:			Dict					=					Arrow(start=lowerCamelCase__   ,				end=lowerCamelCase__   ,				color=lowerCamelCase__   ,				buff=0.5			)
     a.next_to(model_arr[0].get_left()   ,				lowerCamelCase__   ,				buff=0.2			)

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

     __UpperCamelCase	:			int					=					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(lowerCamelCase__   ,				run_time=3			)			)

     __UpperCamelCase	:			Dict					=					{"""run_time""": 1, """fade_in""": True, """fade_out""": True, """buff""": 0.02}

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

     __UpperCamelCase	:			int					=					a.copy()
     for i in range(6			):
        a_c.next_to(model_arr[i].get_right() + 0.02   ,				lowerCamelCase__   ,				buff=0.2			)

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

        __UpperCamelCase	:			Tuple					=					AnimationGroup(
            FadeOut(lowerCamelCase__   ,				run_time=0.5			)   ,				MoveToTarget(lowerCamelCase__   ,				run_time=0.5			)   ,				FadeIn(lowerCamelCase__   ,				run_time=0.5			)   ,				lag_ratio=0.2			)

        self.play(lowerCamelCase__			)

        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	:			Any					=					0.7

           self.play(
               Circumscribe(model_arr[i]   ,				**lowerCamelCase__			)   ,				Circumscribe(cpu_left_col_base[i]   ,				**lowerCamelCase__			)   ,				Circumscribe(cpu_left_col_base[i + 1]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(gpu_rect[0]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(model_arr[i + 1]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				)
           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.02   ,				buff=0.2			)

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

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

        __UpperCamelCase	:			str					=					a_c
        __UpperCamelCase	:			List[Any]					=					a_c.copy()

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

     __UpperCamelCase	:			List[str]					=					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(lowerCamelCase__   ,				run_time=3			)   ,				MoveToTarget(lowerCamelCase__			)			)

     self.wait()



