
import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.generation import (
						    FlaxForcedBOSTokenLogitsProcessor,
						    FlaxForcedEOSTokenLogitsProcessor,
						    FlaxLogitsProcessorList,
						    FlaxMinLengthLogitsProcessor,
						    FlaxTemperatureLogitsWarper,
						    FlaxTopKLogitsWarper,
						    FlaxTopPLogitsWarper,
						)
@require_flax
class 					lowercase__(							unittest.TestCase				):



	"""simple docstring"""




	def 		_lowercase    (			self    :     str			,				SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     int			)			->	Dict:
					lowercase_	    =   jnp.ones((batch_size, length)			) / length
					return scores
	def 		_lowercase    (			self    :     List[Any]			)			->	List[Any]:
					lowercase_	    =   None
					lowercase_	    =   2_0

					lowercase_	    =   self._get_uniform_logits(batch_size=2			,				length=SCREAMING_SNAKE_CASE_			)

					# tweak scores to not be uniform anymore
					lowercase_	    =   scores.at[1, 5].set((1 / length) + 0.1			)  # peak, 1st batch
					lowercase_	    =   scores.at[1, 1_0].set((1 / length) - 0.4			)  # valley, 1st batch

					# compute softmax
					lowercase_	    =   jax.nn.softmax(SCREAMING_SNAKE_CASE_			,				axis=-1			)

					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=0.5			)
					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=1.3			)

					lowercase_	    =   jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE_			,				scores.copy()			,				cur_len=SCREAMING_SNAKE_CASE_			)			,				axis=-1			)
					lowercase_	    =   jax.nn.softmax(temp_dist_warper_smoother(SCREAMING_SNAKE_CASE_			,				scores.copy()			,				cur_len=SCREAMING_SNAKE_CASE_			)			,				axis=-1			)

					# uniform distribution stays uniform
					self.assertTrue(jnp.allclose(probs[0, :]			,				warped_prob_sharp[0, :]			,				atol=1e-3			)			)
					self.assertTrue(jnp.allclose(probs[0, :]			,				warped_prob_smooth[0, :]			,				atol=1e-3			)			)

					# sharp peaks get higher, valleys get lower
					self.assertLess(probs[1, :].max()			,				warped_prob_sharp[1, :].max()			)
					self.assertGreater(probs[1, :].min()			,				warped_prob_sharp[1, :].min()			)

					# smooth peaks get lower, valleys get higher
					self.assertGreater(probs[1, :].max()			,				warped_prob_smooth[1, :].max()			)
					self.assertLess(probs[1, :].min()			,				warped_prob_smooth[1, :].min()			)
	def 		_lowercase    (			self    :     Optional[Any]			)			->	List[str]:
					lowercase_	    =   None
					lowercase_	    =   1_0
					lowercase_	    =   2

					# create ramp distribution
					lowercase_	    =   np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_			)[None, :]			,				(batch_size, vocab_size)			).copy()
					lowercase_	    =   ramp_logits[1:, : vocab_size // 2] + vocab_size

					lowercase_	    =   FlaxTopKLogitsWarper(3			)

					lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

					# check that correct tokens are filtered
					self.assertListEqual(jnp.isinf(scores[0]			).tolist()			,				7 * [True] + 3 * [False]			)
					self.assertListEqual(jnp.isinf(scores[1]			).tolist()			,				2 * [True] + 3 * [False] + 5 * [True]			)

					# check special case
					lowercase_	    =   5
					lowercase_	    =   FlaxTopKLogitsWarper(top_k=1			,				filter_value=0.0			,				min_tokens_to_keep=3			)

					lowercase_	    =   np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_			)[None, :]			,				(batch_size, length)			).copy()
					lowercase_	    =   top_k_warp_safety_check(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

					# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
					self.assertListEqual((scores == 0.0).sum(axis=-1			).tolist()			,				[2, 2]			)
	def 		_lowercase    (			self    :     List[Any]			)			->	Union[str, Any]:
					lowercase_	    =   None
					lowercase_	    =   1_0
					lowercase_	    =   2

					# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
					lowercase_	    =   np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]			)			)

					lowercase_	    =   FlaxTopPLogitsWarper(0.8			)
					lowercase_	    =   np.exp(top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)			)

					# dist should be filtered to keep min num values so that sum is >= top_p
					# exp (-inf) => 0
					lowercase_	    =   np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]			)
					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-3			)			)

					# check edge cases with negative and extreme logits
					lowercase_	    =   np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_			)[None, :]			,				(batch_size, vocab_size)			).copy() - (
					    vocab_size // 2
					)

					# make ramp_logits more extreme
					lowercase_	    =   ramp_logits[1] * 1_00.0

					# make sure at least 2 tokens are kept
					lowercase_	    =   FlaxTopPLogitsWarper(0.9			,				min_tokens_to_keep=2			,				filter_value=0.0			)
					lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

					# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
					self.assertListEqual((filtered_dist != 0.0).sum(axis=-1			).tolist()			,				[3, 2]			)
	def 		_lowercase    (			self    :     Tuple			)			->	Any:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0

					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that min length is applied at length 5
					lowercase_	    =   ids_tensor((batch_size, 2_0)			,				vocab_size=2_0			)
					lowercase_	    =   5
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   min_dist_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()			,				4 * [-float('''inf'''			)]			)

					# check that min length is not applied anymore at length 15
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   1_5
					lowercase_	    =   min_dist_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     List[str]			)			->	Dict:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0

					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the bos_token_id score
					lowercase_	    =   ids_tensor((batch_size, 1)			,				vocab_size=2_0			)
					lowercase_	    =   1
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]			).all()			)
					self.assertListEqual(scores[:, bos_token_id].tolist()			,				4 * [0]			)  # score for bos_token_id shold be zero

					# check that bos_token_id is not forced if current length is greater than 1
					lowercase_	    =   3
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     int			)			->	Optional[int]:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0
					lowercase_	    =   5

					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the eos_token_id when max_length is reached
					lowercase_	    =   ids_tensor((batch_size, 4)			,				vocab_size=2_0			)
					lowercase_	    =   4
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]			).all()			)
					self.assertListEqual(scores[:, eos_token_id].tolist()			,				4 * [0]			)  # score for eos_token_id should be zero

					# check that eos_token_id is not forced if max_length is not reached
					lowercase_	    =   3
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     int			)			->	Tuple:
					lowercase_	    =   4
					lowercase_	    =   1_0
					lowercase_	    =   1_5
					lowercase_	    =   2
					lowercase_	    =   1
					lowercase_	    =   1_5

					# dummy input_ids and scores
					lowercase_	    =   ids_tensor((batch_size, sequence_length)			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   input_ids.copy()

					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   scores.copy()

					# instantiate all dist processors
					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=0.5			)
					lowercase_	    =   FlaxTopKLogitsWarper(3			)
					lowercase_	    =   FlaxTopPLogitsWarper(0.8			)

					# instantiate all logits processors
					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   1_0

					# no processor list
					lowercase_	    =   temp_dist_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   min_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   bos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   eos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

					# with processor list
					lowercase_	    =   FlaxLogitsProcessorList(
					    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
					lowercase_	    =   processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

					# scores should be equal
					self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-3			)			)

					# input_ids should never be changed
					self.assertListEqual(input_ids.tolist()			,				input_ids_comp.tolist()			)




	def 		_lowercase    (			self    :     List[Any]			)			->	List[str]:
					lowercase_	    =   4
					lowercase_	    =   1_0
					lowercase_	    =   1_5
					lowercase_	    =   2
					lowercase_	    =   1
					lowercase_	    =   1_5

					# dummy input_ids and scores
					lowercase_	    =   ids_tensor((batch_size, sequence_length)			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   input_ids.copy()

					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   scores.copy()

					# instantiate all dist processors
					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=0.5			)
					lowercase_	    =   FlaxTopKLogitsWarper(3			)
					lowercase_	    =   FlaxTopPLogitsWarper(0.8			)

					# instantiate all logits processors
					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   1_0

					# no processor list
					def run_no_processor_list(SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     int			):
									lowercase_	    =   temp_dist_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   min_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   bos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   eos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					# with processor list
					def run_processor_list(SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     Dict			,				SCREAMING_SNAKE_CASE_    :     List[Any]			):
									lowercase_	    =   FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
									lowercase_	    =   processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					lowercase_	    =   jax.jit(SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   jax.jit(SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   jitted_run_no_processor_list(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   jitted_run_processor_list(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)

					# scores should be equal
					self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-3			)			)

					# input_ids should never be changed
					self.assertListEqual(input_ids.tolist()			,				input_ids_comp.tolist()			)












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()



