






import torch
from torch import nn





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=1   ,						_SCREAMING_SNAKE_CASE      :    List[str]=False ):






           """simple docstring"""




           super().__init__()

           UpperCamelCase					       =  n_token
           UpperCamelCase					       =  d_embed
           UpperCamelCase					       =  d_proj

           UpperCamelCase					       =  cutoffs + [n_token]
           UpperCamelCase					       =  [0] + self.cutoffs
           UpperCamelCase					       =  div_val

           UpperCamelCase					       =  self.cutoffs[0]
           UpperCamelCase					       =  len(self.cutoffs ) - 1
           UpperCamelCase					       =  self.shortlist_size + self.n_clusters

           if self.n_clusters > 0:
                 UpperCamelCase					       =  nn.Parameter(torch.zeros(self.n_clusters   ,						self.d_embed ) )
                 UpperCamelCase					       =  nn.Parameter(torch.zeros(self.n_clusters ) )

           UpperCamelCase					       =  nn.ModuleList()
           UpperCamelCase					       =  nn.ParameterList()

           if div_val == 1:
                 for i in range(len(self.cutoffs ) ):
                       if d_proj != d_embed:
                             self.out_projs.append(nn.Parameter(torch.FloatTensor(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) ) )
                       else:
                             self.out_projs.append(_SCREAMING_SNAKE_CASE )

                 self.out_layers.append(nn.Linear(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) )
           else:
                 for i in range(len(self.cutoffs ) ):
                       UpperCamelCase       ,  UpperCamelCase					       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                       UpperCamelCase					       =  d_embed // (div_val**i)

                       self.out_projs.append(nn.Parameter(torch.FloatTensor(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) ) )

                       self.out_layers.append(nn.Linear(_SCREAMING_SNAKE_CASE   ,						r_idx - l_idx ) )

           UpperCamelCase					       =  keep_order







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






           """simple docstring"""




           if proj is None:
                 UpperCamelCase					       =  nn.functional.linear(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						bias=_SCREAMING_SNAKE_CASE )
           else:
                 # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
                 UpperCamelCase					       =  nn.functional.linear(_SCREAMING_SNAKE_CASE   ,						proj.t().contiguous() )
                 UpperCamelCase					       =  nn.functional.linear(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						bias=_SCREAMING_SNAKE_CASE )
                 # else:
                 #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
                 #     if bias is not None:
                 #         logit = logit + bias

           return logit







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=False ):






           """simple docstring"""




           if labels is not None:
                 # Shift so that tokens < n predict n
                 UpperCamelCase					       =  hidden[..., :-1, :].contiguous()
                 UpperCamelCase					       =  labels[..., 1:].contiguous()
                 UpperCamelCase					       =  hidden.view(-1   ,						hidden.size(-1 ) )
                 UpperCamelCase					       =  labels.view(-1 )
                 if hidden.size(0 ) != labels.size(0 ):
                       raise RuntimeError('Input and labels should have the same size in the batch dimension.' )
           else:
                 UpperCamelCase					       =  hidden.view(-1   ,						hidden.size(-1 ) )

           if self.n_clusters == 0:
                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						self.out_layers[0].weight   ,						self.out_layers[0].bias   ,						self.out_projs[0] )
                 if labels is not None:
                       UpperCamelCase					       =  labels != -100
                       UpperCamelCase					       =  torch.zeros_like(_SCREAMING_SNAKE_CASE   ,						dtype=hidden.dtype   ,						device=hidden.device )
                       UpperCamelCase					       =  (
                           -nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=-1 )[mask].gather(1   ,						labels[mask].unsqueeze(1 ) ).squeeze(1 )
                       )
                 else:
                       UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=-1 )
           else:
                 # construct weights and biases
                 UpperCamelCase       ,  UpperCamelCase					       =  [], []
                 for i in range(len(self.cutoffs ) ):
                       if self.div_val == 1:
                             UpperCamelCase       ,  UpperCamelCase					       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             UpperCamelCase					       =  self.out_layers[0].weight[l_idx:r_idx]
                             UpperCamelCase					       =  self.out_layers[0].bias[l_idx:r_idx]
                       else:
                             UpperCamelCase					       =  self.out_layers[i].weight
                             UpperCamelCase					       =  self.out_layers[i].bias

                       if i == 0:
                             UpperCamelCase					       =  torch.cat([weight_i, self.cluster_weight]   ,						dim=0 )
                             UpperCamelCase					       =  torch.cat([bias_i, self.cluster_bias]   ,						dim=0 )

                       weights.append(_SCREAMING_SNAKE_CASE )
                       biases.append(_SCREAMING_SNAKE_CASE )

                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[0], biases[0], self.out_projs[0]

                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

                 if labels is None:
                       UpperCamelCase					       =  hidden.new_empty((head_logit.size(0 ), self.n_token) )
                 else:
                       UpperCamelCase					       =  torch.zeros_like(_SCREAMING_SNAKE_CASE   ,						dtype=hidden.dtype   ,						device=hidden.device )

                 UpperCamelCase					       =  0
                 UpperCamelCase					       =  [0] + self.cutoffs
                 for i in range(len(_SCREAMING_SNAKE_CASE ) - 1 ):
                       UpperCamelCase       ,  UpperCamelCase					       =  cutoff_values[i], cutoff_values[i + 1]

                       if labels is not None:
                             UpperCamelCase					       =  (labels >= l_idx) & (labels < r_idx)
                             UpperCamelCase					       =  mask_i.nonzero().squeeze()

                             if indices_i.numel() == 0:
                                   continue

                             UpperCamelCase					       =  labels.index_select(0   ,						_SCREAMING_SNAKE_CASE ) - l_idx
                             UpperCamelCase					       =  head_logprob.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  hidden.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                       else:
                             UpperCamelCase					       =  hidden

                       if i == 0:
                             if labels is not None:
                                   UpperCamelCase					       =  head_logprob_i.gather(1   ,						target_i[:, None] ).squeeze(1 )
                             else:
                                   UpperCamelCase					       =  head_logprob[:, : self.cutoffs[0]]
                       else:
                             UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[i], biases[i], self.out_projs[i]

                             UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )
                             UpperCamelCase					       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
                             if labels is not None:
                                   UpperCamelCase					       =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                                       1   ,						target_i[:, None] ).squeeze(1 )
                             else:
                                   UpperCamelCase					       =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                                   UpperCamelCase					       =  logprob_i

                       if labels is not None:
                             if (hasattr(self   ,						'keep_order' ) and self.keep_order) or keep_order:
                                   out.index_copy_(0   ,						_SCREAMING_SNAKE_CASE   ,						-logprob_i )
                             else:
                                   out[offset : offset + logprob_i.size(0 )].copy_(-logprob_i )
                             offset += logprob_i.size(0 )

           return out







     def   _SCREAMING_SNAKE_CASE     ( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           if self.n_clusters == 0:
                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						self.out_layers[0].weight   ,						self.out_layers[0].bias   ,						self.out_projs[0] )
                 return nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=-1 )
           else:
                 # construct weights and biases
                 UpperCamelCase       ,  UpperCamelCase					       =  [], []
                 for i in range(len(self.cutoffs ) ):
                       if self.div_val == 1:
                             UpperCamelCase       ,  UpperCamelCase					       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             UpperCamelCase					       =  self.out_layers[0].weight[l_idx:r_idx]
                             UpperCamelCase					       =  self.out_layers[0].bias[l_idx:r_idx]
                       else:
                             UpperCamelCase					       =  self.out_layers[i].weight
                             UpperCamelCase					       =  self.out_layers[i].bias

                       if i == 0:
                             UpperCamelCase					       =  torch.cat([weight_i, self.cluster_weight]   ,						dim=0 )
                             UpperCamelCase					       =  torch.cat([bias_i, self.cluster_bias]   ,						dim=0 )

                       weights.append(_SCREAMING_SNAKE_CASE )
                       biases.append(_SCREAMING_SNAKE_CASE )

                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[0], biases[0], self.out_projs[0]
                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

                 UpperCamelCase					       =  hidden.new_empty((head_logit.size(0 ), self.n_token) )
                 UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

                 UpperCamelCase					       =  [0] + self.cutoffs
                 for i in range(len(_SCREAMING_SNAKE_CASE ) - 1 ):
                       UpperCamelCase       ,  UpperCamelCase					       =  cutoff_values[i], cutoff_values[i + 1]

                       if i == 0:
                             UpperCamelCase					       =  head_logprob[:, : self.cutoffs[0]]
                       else:
                             UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[i], biases[i], self.out_projs[i]

                             UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

                             UpperCamelCase					       =  head_logprob[:, -i] + tail_logprob_i
                             UpperCamelCase					       =  logprob_i

                 return out











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



