





"""simple docstring"""






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
    import os

    # The slow tests are often failing with OOM error on GPU
    # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
    # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    lowerCAmelCase_      =				'''platform'''

    import jax
    import jax.numpy as jnp

    from transformers import AutoTokenizer
    from transformers.models.mbart.modeling_flax_mbart import (
        FlaxMBartForConditionalGeneration,
        FlaxMBartForQuestionAnswering,
        FlaxMBartForSequenceClassification,
        FlaxMBartModel,
        shift_tokens_right,
    )

def    __lowerCamelCase     (							SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE=None,SCREAMING_SNAKE_CASE=None,SCREAMING_SNAKE_CASE=None,SCREAMING_SNAKE_CASE=None,SCREAMING_SNAKE_CASE=None,SCREAMING_SNAKE_CASE=None,)		->							str:






       """simple docstring"""




       if attention_mask is None:
              _UpperCAmelCase			     =  np.where(input_ids != config.pad_token_id,1,0				)
       if decoder_attention_mask is None:
              _UpperCAmelCase			     =  np.where(decoder_input_ids != config.pad_token_id,1,0				)
       if head_mask is None:
              _UpperCAmelCase			     =  np.ones((config.encoder_layers, config.encoder_attention_heads)				)
       if decoder_head_mask is None:
              _UpperCAmelCase			     =  np.ones((config.decoder_layers, config.decoder_attention_heads)				)
       if cross_attn_head_mask is None:
              _UpperCAmelCase			     =  np.ones((config.decoder_layers, config.decoder_attention_heads)				)
       return {
           "input_ids": input_ids,
           "decoder_input_ids": decoder_input_ids,
           "attention_mask": attention_mask,
           "decoder_attention_mask": decoder_attention_mask,
       }



class 						lowerCAmelCase       (   unittest.TestCase				):






       def __init__( self     ,	a__     ,	a__=13     ,	a__=7     ,	a__=True     ,	a__=False     ,	a__=99     ,	a__=16     ,	a__=2     ,	a__=4     ,	a__=4     ,	a__="gelu"     ,	a__=0.1     ,	a__=0.1     ,	a__=32     ,	a__=2     ,	a__=1     ,	a__=0     ,	a__=2     ,	a__=0.02     ,	):
              _UpperCAmelCase			     =  parent
              _UpperCAmelCase			     =  batch_size
              _UpperCAmelCase			     =  seq_length
              _UpperCAmelCase			     =  is_training
              _UpperCAmelCase			     =  use_labels
              _UpperCAmelCase			     =  vocab_size
              _UpperCAmelCase			     =  hidden_size
              _UpperCAmelCase			     =  num_hidden_layers
              _UpperCAmelCase			     =  num_attention_heads
              _UpperCAmelCase			     =  intermediate_size
              _UpperCAmelCase			     =  hidden_act
              _UpperCAmelCase			     =  hidden_dropout_prob
              _UpperCAmelCase			     =  attention_probs_dropout_prob
              _UpperCAmelCase			     =  max_position_embeddings
              _UpperCAmelCase			     =  eos_token_id
              _UpperCAmelCase			     =  pad_token_id
              _UpperCAmelCase			     =  bos_token_id
              _UpperCAmelCase			     =  decoder_start_token_id
              _UpperCAmelCase			     =  initializer_range






       def 				__A			( self  ):
              _UpperCAmelCase			     =  np.clip(ids_tensor([self.batch_size, self.seq_length - 1]     ,	self.vocab_size  )     ,	3     ,	self.vocab_size  )
              _UpperCAmelCase			     =  np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)     ,	dtype=np.intaa  ))     ,	-1  )

              _UpperCAmelCase			     =  shift_tokens_right(a__     ,	1  )

              _UpperCAmelCase			     =  MBartConfig(
                  vocab_size=self.vocab_size     ,	d_model=self.hidden_size     ,	encoder_layers=self.num_hidden_layers     ,	decoder_layers=self.num_hidden_layers     ,	encoder_attention_heads=self.num_attention_heads     ,	decoder_attention_heads=self.num_attention_heads     ,	encoder_ffn_dim=self.intermediate_size     ,	decoder_ffn_dim=self.intermediate_size     ,	dropout=self.hidden_dropout_prob     ,	attention_dropout=self.attention_probs_dropout_prob     ,	max_position_embeddings=self.max_position_embeddings     ,	eos_token_id=self.eos_token_id     ,	bos_token_id=self.bos_token_id     ,	pad_token_id=self.pad_token_id     ,	decoder_start_token_id=self.decoder_start_token_id     ,	initializer_range=self.initializer_range     ,	use_cache=a__     ,	)
              _UpperCAmelCase			     =  prepare_mbart_inputs_dict(a__     ,	a__     ,	a__  )
              return config, inputs_dict






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.prepare_config_and_inputs()
              return config, inputs_dict






       def 				__A			( self     ,	a__     ,	a__     ,	a__  ):
              _UpperCAmelCase			     =  20
              _UpperCAmelCase			     =  model_class_name(a__  )

              _UpperCAmelCase			     =  model.encode(inputs_dict['input_ids']  )

              _UpperCAmelCase     ,				_UpperCAmelCase			     =  (
                  inputs_dict['decoder_input_ids'],
                  inputs_dict['decoder_attention_mask'],
              )

              _UpperCAmelCase			     =  model.init_cache(decoder_input_ids.shape[0]     ,	a__     ,	a__  )
              _UpperCAmelCase			     =  jnp.ones((decoder_input_ids.shape[0], max_decoder_length)     ,	dtype='i4'  )

              _UpperCAmelCase			     =  jnp.broadcast_to(
                  jnp.arange(decoder_input_ids.shape[-1] - 1  )[None, :]     ,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)     ,	)
              _UpperCAmelCase			     =  model.decode(
                  decoder_input_ids[:, :-1]     ,	a__     ,	decoder_attention_mask=a__     ,	past_key_values=a__     ,	decoder_position_ids=a__     ,	)

              _UpperCAmelCase			     =  jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]     ,	dtype='i4'  )
              _UpperCAmelCase			     =  model.decode(
                  decoder_input_ids[:, -1:]     ,	a__     ,	decoder_attention_mask=a__     ,	past_key_values=outputs_cache.past_key_values     ,	decoder_position_ids=a__     ,	)

              _UpperCAmelCase			     =  model.decode(a__     ,	a__  )

              _UpperCAmelCase			     =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
              self.parent.assertTrue(diff < 1E-3     ,	msg=f"""Max diff is {diff}"""  )






       def 				__A			( self     ,	a__     ,	a__     ,	a__  ):
              _UpperCAmelCase			     =  20
              _UpperCAmelCase			     =  model_class_name(a__  )

              _UpperCAmelCase			     =  model.encode(inputs_dict['input_ids']  )

              _UpperCAmelCase     ,				_UpperCAmelCase			     =  (
                  inputs_dict['decoder_input_ids'],
                  inputs_dict['decoder_attention_mask'],
              )

              _UpperCAmelCase			     =  jnp.concatenate(
                  [
                      decoder_attention_mask,
                      jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])  ),
                  ]     ,	axis=-1     ,	)

              _UpperCAmelCase			     =  model.init_cache(decoder_input_ids.shape[0]     ,	a__     ,	a__  )
              _UpperCAmelCase			     =  jnp.broadcast_to(
                  jnp.arange(decoder_input_ids.shape[-1] - 1  )[None, :]     ,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)     ,	)

              _UpperCAmelCase			     =  model.decode(
                  decoder_input_ids[:, :-1]     ,	a__     ,	decoder_attention_mask=a__     ,	past_key_values=a__     ,	decoder_position_ids=a__     ,	)
              _UpperCAmelCase			     =  jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]     ,	dtype='i4'  )
              _UpperCAmelCase			     =  model.decode(
                  decoder_input_ids[:, -1:]     ,	a__     ,	past_key_values=outputs_cache.past_key_values     ,	decoder_attention_mask=a__     ,	decoder_position_ids=a__     ,	)

              _UpperCAmelCase			     =  model.decode(a__     ,	a__     ,	decoder_attention_mask=a__  )

              _UpperCAmelCase			     =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
              self.parent.assertTrue(diff < 1E-3     ,	msg=f"""Max diff is {diff}"""  )



@require_flax
class 						lowerCAmelCase       (   unittest.TestCase				):
       lowerCAmelCase__    			=	99






       def 				__A			( self  ):
              _UpperCAmelCase			     =  np.array(
                  [
                      [71, 82, 18, 33, 46, 91, 2],
                      [68, 34, 26, 58, 30, 82, 2],
                      [5, 97, 17, 39, 94, 40, 2],
                      [76, 83, 94, 25, 70, 78, 2],
                      [87, 59, 41, 35, 48, 66, 2],
                      [55, 13, 16, 58, 5, 2, 1],  # note padding
                      [64, 27, 31, 51, 12, 75, 2],
                      [52, 64, 86, 17, 83, 39, 2],
                      [48, 61, 9, 24, 71, 82, 2],
                      [26, 1, 60, 48, 22, 13, 2],
                      [21, 5, 62, 28, 14, 76, 2],
                      [45, 98, 37, 86, 59, 48, 2],
                      [70, 70, 50, 9, 28, 0, 2],
                  ]     ,	dtype=np.intaa     ,	)

              _UpperCAmelCase			     =  input_ids.shape[0]
              _UpperCAmelCase			     =  MBartConfig(
                  vocab_size=self.vocab_size     ,	d_model=24     ,	encoder_layers=2     ,	decoder_layers=2     ,	encoder_attention_heads=2     ,	decoder_attention_heads=2     ,	encoder_ffn_dim=32     ,	decoder_ffn_dim=32     ,	max_position_embeddings=48     ,	eos_token_id=2     ,	pad_token_id=1     ,	bos_token_id=0     ,	)
              return config, input_ids, batch_size






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase			     =  self._get_config_and_data()
              _UpperCAmelCase			     =  FlaxMBartForSequenceClassification(a__  )
              _UpperCAmelCase			     =  model(input_ids=a__     ,	decoder_input_ids=a__  )
              _UpperCAmelCase			     =  (batch_size, config.num_labels)
              self.assertEqual(outputs['logits'].shape     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase			     =  self._get_config_and_data()
              _UpperCAmelCase			     =  FlaxMBartForQuestionAnswering(a__  )
              _UpperCAmelCase			     =  model(input_ids=a__  )

              self.assertEqual(outputs['start_logits'].shape     ,	input_ids.shape  )
              self.assertEqual(outputs['end_logits'].shape     ,	input_ids.shape  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase			     =  self._get_config_and_data()
              _UpperCAmelCase			     =  FlaxMBartForConditionalGeneration(a__  )
              _UpperCAmelCase			     =  lm_model(input_ids=a__  )
              _UpperCAmelCase			     =  (batch_size, input_ids.shape[1], config.vocab_size)
              self.assertEqual(outputs['logits'].shape     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  MBartConfig(
                  vocab_size=self.vocab_size     ,	d_model=14     ,	encoder_layers=2     ,	decoder_layers=2     ,	encoder_attention_heads=2     ,	decoder_attention_heads=2     ,	encoder_ffn_dim=8     ,	decoder_ffn_dim=8     ,	max_position_embeddings=48     ,	)
              _UpperCAmelCase			     =  FlaxMBartForConditionalGeneration(a__  )
              _UpperCAmelCase			     =  np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]     ,	dtype=np.intaa  )
              _UpperCAmelCase			     =  np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]     ,	dtype=np.intaa  )
              _UpperCAmelCase			     =  lm_model(input_ids=a__     ,	decoder_input_ids=a__  )
              _UpperCAmelCase			     =  (*summary.shape, config.vocab_size)
              self.assertEqual(outputs['logits'].shape     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]     ,	dtype=np.intaa  )
              _UpperCAmelCase			     =  shift_tokens_right(a__     ,	1  )
              _UpperCAmelCase			     =  np.equal(a__     ,	1  ).astype(np.floataa  ).sum()
              _UpperCAmelCase			     =  np.equal(a__     ,	1  ).astype(np.floataa  ).sum()
              self.assertEqual(shifted.shape     ,	input_ids.shape  )
              self.assertEqual(a__     ,	n_pad_before - 1  )
              self.assertTrue(np.equal(shifted[:, 0]     ,	2  ).all()  )



@require_flax
class 						lowerCAmelCase       (   __snake_case      ,						unittest.TestCase      ,						__snake_case				):
       lowerCAmelCase__    			=	True
       lowerCAmelCase__    			=	(
           (
               FlaxMBartModel,
               FlaxMBartForConditionalGeneration,
               FlaxMBartForSequenceClassification,
               FlaxMBartForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )
       lowerCAmelCase__    			=	(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()






       def 				__A			( self  ):
              _UpperCAmelCase			     =  FlaxMBartModelTester(self  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              for model_class in self.all_model_classes:
                     self.model_tester.check_use_cache_forward(a__     ,	a__     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              for model_class in self.all_model_classes:
                     self.model_tester.check_use_cache_forward_with_attn_mask(a__     ,	a__     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     with self.subTest(model_class.__name__  ):
                            _UpperCAmelCase			     =  self._prepare_for_class(a__     ,	a__  )
                            _UpperCAmelCase			     =  model_class(a__  )

                            @jax.jit
                            def encode_jitted(a__     ,	a__=None     ,	**a__  ):
                                   return model.encode(input_ids=a__     ,	attention_mask=a__  )

                            with self.subTest('JIT Enabled'  ):
                                   _UpperCAmelCase			     =  encode_jitted(**a__  ).to_tuple()

                            with self.subTest('JIT Disabled'  ):
                                   with jax.disable_jit():
                                          _UpperCAmelCase			     =  encode_jitted(**a__  ).to_tuple()

                            self.assertEqual(len(a__  )     ,	len(a__  )  )
                            for jitted_output, output in zip(a__     ,	a__  ):
                                   self.assertEqual(jitted_output.shape     ,	output.shape  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     with self.subTest(model_class.__name__  ):
                            _UpperCAmelCase			     =  model_class(a__  )
                            _UpperCAmelCase			     =  model.encode(inputs_dict['input_ids']     ,	inputs_dict['attention_mask']  )

                            _UpperCAmelCase			     =  {
                                'decoder_input_ids': inputs_dict['decoder_input_ids'],
                                'decoder_attention_mask': inputs_dict['decoder_attention_mask'],
                                'encoder_outputs': encoder_outputs,
                            }

                            @jax.jit
                            def decode_jitted(a__     ,	a__     ,	a__  ):
                                   return model.decode(
                                       decoder_input_ids=a__     ,	decoder_attention_mask=a__     ,	encoder_outputs=a__     ,	)

                            with self.subTest('JIT Enabled'  ):
                                   _UpperCAmelCase			     =  decode_jitted(**a__  ).to_tuple()

                            with self.subTest('JIT Disabled'  ):
                                   with jax.disable_jit():
                                          _UpperCAmelCase			     =  decode_jitted(**a__  ).to_tuple()

                            self.assertEqual(len(a__  )     ,	len(a__  )  )
                            for jitted_output, output in zip(a__     ,	a__  ):
                                   self.assertEqual(jitted_output.shape     ,	output.shape  )






       @slow
       def 				__A			( self  ):
              for model_class_name in self.all_model_classes:
                     _UpperCAmelCase			     =  model_class_name.from_pretrained('facebook/mbart-large-cc25'     ,	from_pt=a__  )
                     # FlaxMBartForSequenceClassification expects eos token in input_ids
                     _UpperCAmelCase			     =  np.ones((1, 1)  ) * model.config.eos_token_id
                     _UpperCAmelCase			     =  model(a__  )
                     self.assertIsNotNone(a__  )



@require_flax
@require_sentencepiece
@require_tokenizers
class 						lowerCAmelCase       (   unittest.TestCase				):
       lowerCAmelCase__    			=	[
           ''' UN Chief Says There Is No Military Solution in Syria''',
       ]
       lowerCAmelCase__    			=	[
           '''Şeful ONU declară că nu există o soluţie militară în Siria''',
       ]
       lowerCAmelCase__    			=	'''facebook/mbart-large-en-ro'''






       @cached_property
       def 				__A			( self  ):
              return AutoTokenizer.from_pretrained(self.model_name  )






       @cached_property
       def 				__A			( self  ):
              _UpperCAmelCase			     =  FlaxMBartForConditionalGeneration.from_pretrained(self.model_name     ,	from_pt=a__  )
              return model






       def 				__A			( self     ,	**a__  ):
              _UpperCAmelCase			     =  self.translate_src_text(**a__  )
              self.assertListEqual(self.expected_text     ,	a__  )






       def 				__A			( self     ,	**a__  ):
              _UpperCAmelCase			     =  self.tokenizer(self.src_text     ,	**a__     ,	return_tensors='np'  )
              _UpperCAmelCase			     =  self.model.generate(
                  model_inputs.input_ids     ,	attention_mask=model_inputs.attention_mask     ,	decoder_start_token_id=self.tokenizer.lang_code_to_id['ro_RO']     ,	early_stopping=a__     ,	num_beams=2     ,	).sequences
              _UpperCAmelCase			     =  self.tokenizer.batch_decode(a__     ,	skip_special_tokens=a__  )
              return generated_words






       @slow
       def 				__A			( self  ):
              self._assert_generated_batch_equal_expected()









"""simple docstring"""






from manim import *
class 	_lowerCAmelCase    (      __snake_case       ):

							def 					_lowerCAmelCase  ( self       :    Dict    )   ->    Tuple:

									"""simple docstring"""
									lowercase										=							Rectangle(height=0.5 ,		width=0.5    )
									lowercase										=							Rectangle(height=0.46 ,		width=0.46    ).set_stroke(width=0    )
									lowercase										=							Rectangle(height=0.25 ,		width=0.25    )

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

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

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

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

									lowercase										=							[]
									lowercase										=							[]

									for i, rect in enumerate(a    ):
											rect.set_stroke(a    )

											lowercase										=							Rectangle(height=0.46 / 4 ,		width=0.46 / 3    ).set_stroke(width=0.0    ).set_fill(a ,		opacity=0.7    )

											if i == 0:
													cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT    ) ,		buff=0.02 ,		direction=a    )
													cpu_target.set_x(cpu_target.get_x() + 0.1    )
											elif i == 3:
													cpu_target.next_to(model_cpu_arr[0] ,		direction=a ,		buff=0.0    )
											else:
													cpu_target.next_to(model_cpu_arr[i - 1] ,		direction=a ,		buff=0.0    )
											self.add(a    )
											model_cpu_arr.append(a    )

									self.add(*a ,		*a    )

									lowercase										=							[meta_mem.copy() for i in range(6    )]
									lowercase										=							[meta_mem.copy() for i in range(6    )]
									lowercase										=							VGroup(*a    ).arrange(a ,		buff=0    )
									lowercase										=							VGroup(*a    ).arrange(a ,		buff=0    )
									lowercase										=							VGroup(a ,		a    ).arrange(a ,		buff=0    )
									lowercase										=							Text('''Disk''' ,		font_size=24    )
									lowercase										=							Group(a ,		a    ).arrange(a ,		buff=0.5 ,		aligned_edge=a    )
									disk.move_to([-4.0, -1.25, 0]    )
									self.add(a ,		a    )

									lowercase										=							[]

									for i in range(6    ):
											lowercase										=							fill.copy().set_fill(a ,		opacity=0.8    )
											target.move_to(disk_left_col_base[i]    ).scale(0.5    )
											cpu_disk_arr.append(a    )

									self.add(*a    )

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

									lowercase										=							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(a ,		a    )

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

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

									lowercase										=							MarkupText(f"""The offloaded weights are all sent to the CPU.""" ,		font_size=24    )
									step_a.move_to([2, 2, 0]    )

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

									for i in range(6    ):
											lowercase										=							cpu_disk_arr[i]
											lowercase										=							rect.copy().set_fill(a ,		opacity=0.8    ).scale(2.0    )
											cpa.generate_target()
											cpa.target.move_to(model_base[i]    )

											if i == 0:
													rect.set_fill(a ,		opacity=0.8    )
													rect.generate_target()
													rect.target.move_to(cpu_left_col_base[0]    ).scale(2.0    )

													self.remove(
													    *a ,		*a ,		)

											else:
													rect.generate_target()
													rect.target.move_to(cpu_left_col_base[i]    ).scale(2.0    )
											self.play(MoveToTarget(a    ) ,		MoveToTarget(a    ) ,		model_base[i].animate.set_stroke(a    )    )
									self.play(FadeOut(a    )    )

									lowercase										=							MarkupText(
									    f"""Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed.""" ,		font_size=24 ,		)
									step_a.move_to([2, 2, 0]    )

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

									lowercase										=							[]
									lowercase										=							[]
									for i in range(6    ):
											lowercase										=							Arrow(start=a ,		end=a ,		color=a ,		buff=0.5    )
											a.next_to(model_base[i].get_left() ,		a ,		buff=0.2    )
											arrows.append(a    )
											animations.append(Write(a    )    )
									self.play(*a    )
									self.wait()