




from manim import *







class     __lowerCAmelCase (       __magic_name__ ):

			"""simple docstring"""







			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    Rectangle(height=0.5			,						width=0.5							)
							__lowerCamelCase		       =    Rectangle(height=0.46			,						width=0.46							).set_stroke(width=0							)

							__lowerCamelCase		       =    Rectangle(height=0.25			,						width=0.25							)

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

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

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

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

							__lowerCamelCase		       =    []
							__lowerCamelCase		       =    []

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

											__lowerCamelCase		       =    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__							)

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

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

							__lowerCamelCase		       =    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__							)

							__lowerCamelCase		       =    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__							)

							__lowerCamelCase		       =    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__							)							)

							__lowerCamelCase		       =    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__							)							)

							__lowerCamelCase		       =    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]							)

							__lowerCamelCase		       =    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							)							)

							__lowerCamelCase		       =    {'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]							)							)

							__lowerCamelCase		       =    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							)

											__lowerCamelCase		       =    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:
																			__lowerCamelCase		       =    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]							)							)

											__lowerCamelCase		       =    a_c
											__lowerCamelCase		       =    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							)			,						)

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








from dataclasses import dataclass
from typing import Tuple

import numpy as np
import torch







@dataclass
class     __lowerCAmelCase :

			"""simple docstring"""







			snake_case_   		=       42  # [batch_size x 3]
			snake_case_   		=       42  # [batch_size x 3]
			snake_case_   		=       42  # [batch_size x 3]
			snake_case_   		=       42  # [batch_size x 3]
			snake_case_   		=       42
			snake_case_   		=       42
			snake_case_   		=       42
			snake_case_   		=       42
			snake_case_   		=       42


			def       lowercase_    (      self							)    ->      Union[str, Any]:



							'''simple docstring'''
							assert self.x.shape[0] == self.y.shape[0] == self.z.shape[0] == self.origin.shape[0]
							assert self.x.shape[1] == self.y.shape[1] == self.z.shape[1] == self.origin.shape[1] == 3
							assert len(self.x.shape							) == len(self.y.shape							) == len(self.z.shape							) == len(self.origin.shape							) == 2


			def       lowercase_    (      self							)    ->      Optional[int]:



							'''simple docstring'''
							return torch.from_numpy(np.array([self.width, self.height]			,						dtype=np.floataa							)							)


			def       lowercase_    (      self							)    ->      Optional[Any]:



							'''simple docstring'''
							return torch.from_numpy(np.array([self.x_fov, self.y_fov]			,						dtype=np.floataa							)							)


			def       lowercase_    (      self							)    ->      torch.Tensor:



							'''simple docstring'''
							__lowerCamelCase		       =    torch.arange(self.height * self.width							)
							__lowerCamelCase		       =    torch.stack(
							    [
							        pixel_indices % self.width,
							        torch.div(lowerCamelCase__			,						self.width			,						rounding_mode='trunc'							),
							    ]			,						axis=1			,						)
							return coords


			@property
			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase       ,						*__lowerCamelCase		       =    self.shape
							__lowerCamelCase		       =    int(np.prod(lowerCamelCase__							)							)

							__lowerCamelCase		       =    self.get_image_coords()
							__lowerCamelCase		       =    torch.broadcast_to(coords.unsqueeze(0							)			,						[batch_size * inner_batch_size, *coords.shape]							)
							__lowerCamelCase		       =    self.get_camera_rays(lowerCamelCase__							)

							__lowerCamelCase		       =    rays.view(lowerCamelCase__			,						inner_batch_size * self.height * self.width			,						2			,						3							)

							return rays


			def       lowercase_    (      self			,						lowerCamelCase__							)    ->      torch.Tensor:



							'''simple docstring'''
							__lowerCamelCase       ,						*__lowerCamelCase       ,						__lowerCamelCase		       =    coords.shape
							assert n_coords == 2
							assert batch_size == self.origin.shape[0]

							__lowerCamelCase		       =    coords.view(lowerCamelCase__			,						-1			,						2							)

							__lowerCamelCase		       =    self.resolution()
							__lowerCamelCase		       =    self.fov()

							__lowerCamelCase		       =    (flat.float() / (res - 1)) * 2 - 1
							__lowerCamelCase		       =    fracs * torch.tan(fov / 2							)

							__lowerCamelCase		       =    fracs.view(lowerCamelCase__			,						-1			,						2							)
							__lowerCamelCase		       =    (
							    self.z.view(lowerCamelCase__			,						1			,						3							)
							    + self.x.view(lowerCamelCase__			,						1			,						3							) * fracs[:, :, :1]
							    + self.y.view(lowerCamelCase__			,						1			,						3							) * fracs[:, :, 1:]
							)
							__lowerCamelCase		       =    directions / directions.norm(dim=-1			,						keepdim=lowerCamelCase__							)
							__lowerCamelCase		       =    torch.stack(
							    [
							        torch.broadcast_to(self.origin.view(lowerCamelCase__			,						1			,						3							)			,						[batch_size, directions.shape[1], 3]							),
							        directions,
							    ]			,						dim=2			,						)
							return rays.view(lowerCamelCase__			,						*lowerCamelCase__			,						2			,						3							)






			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      "DifferentiableProjectiveCamera":



							'''simple docstring'''
							assert width * self.height == height * self.width, "The aspect ratio should not change."
							return DifferentiableProjectiveCamera(
							    origin=self.origin			,						x=self.x			,						y=self.y			,						z=self.z			,						width=lowerCamelCase__			,						height=lowerCamelCase__			,						x_fov=self.x_fov			,						y_fov=self.y_fov			,						)







def 			lowerCamelCase_		(       UpperCamelCase__				:      int      )     ->       DifferentiableProjectiveCamera:





				"""simple docstring"""


				__lowerCamelCase		       =    []
				__lowerCamelCase		       =    []
				__lowerCamelCase		       =    []
				__lowerCamelCase		       =    []
				for theta in np.linspace(0 ,     2 * np.pi ,     num=20      ):
								__lowerCamelCase		       =    np.array([np.sin(UpperCamelCase__      ), np.cos(UpperCamelCase__      ), -0.5]      )
								z /= np.sqrt(np.sum(z**2      )      )
								__lowerCamelCase		       =    -z * 4
								__lowerCamelCase		       =    np.array([np.cos(UpperCamelCase__      ), -np.sin(UpperCamelCase__      ), 0.0]      )
								__lowerCamelCase		       =    np.cross(UpperCamelCase__ ,     UpperCamelCase__      )
								origins.append(UpperCamelCase__      )
								xs.append(UpperCamelCase__      )
								ys.append(UpperCamelCase__      )
								zs.append(UpperCamelCase__      )
				return DifferentiableProjectiveCamera(
				    origin=torch.from_numpy(np.stack(UpperCamelCase__ ,     axis=0      )      ).float() ,     x=torch.from_numpy(np.stack(UpperCamelCase__ ,     axis=0      )      ).float() ,     y=torch.from_numpy(np.stack(UpperCamelCase__ ,     axis=0      )      ).float() ,     z=torch.from_numpy(np.stack(UpperCamelCase__ ,     axis=0      )      ).float() ,     width=UpperCamelCase__ ,     height=UpperCamelCase__ ,     x_fov=0.7 ,     y_fov=0.7 ,     shape=(1, len(UpperCamelCase__      )) ,     )


