
"""simple docstring"""





import tensorflow as tf

from ...tf_utils import shape_list



class      a__				(	tf.keras.layers.Layer   ):


			def __init__(      self      :    Union[str, Any]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    List[str]=1      ,  UpperCamelCase_      :    Tuple=False      ,  **UpperCamelCase_      :    Union[str, Any]):


						"""simple docstring"""




						super().__init__(**UpperCamelCase_)

						__UpperCAmelCase	: int		=					vocab_size
						__UpperCAmelCase	: str		=					d_embed
						__UpperCAmelCase	: Dict		=					d_proj

						__UpperCAmelCase	: Any		=					cutoffs + [vocab_size]
						__UpperCAmelCase	: Dict		=					[0] + self.cutoffs
						__UpperCAmelCase	: Optional[int]		=					div_val

						__UpperCAmelCase	: str		=					self.cutoffs[0]
						__UpperCAmelCase	: Dict		=					len(self.cutoffs) - 1
						__UpperCAmelCase	: List[Any]		=					self.shortlist_size + self.n_clusters
						__UpperCAmelCase	: List[str]		=					keep_order

						__UpperCAmelCase	: str		=					[]
						__UpperCAmelCase	: str		=					[]


			def       a_     (      self      :    int      ,  UpperCamelCase_      :    List[str]):


						"""simple docstring"""




						if self.n_clusters > 0:
									__UpperCAmelCase	: Any		=					self.add_weight(
									    shape=(self.n_clusters, self.d_embed)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name="cluster_weight")
									__UpperCAmelCase	: List[Any]		=					self.add_weight(
									    shape=(self.n_clusters,)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name="cluster_bias")

						if self.div_val == 1:
									for i in range(len(self.cutoffs)):
												if self.d_proj != self.d_embed:
															__UpperCAmelCase	: Dict		=					self.add_weight(
															    shape=(self.d_embed, self.d_proj)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_projs_._{i}"      ,  )
															self.out_projs.append(UpperCamelCase_)
												else:
															self.out_projs.append(UpperCamelCase_)
												__UpperCAmelCase	: List[str]		=					self.add_weight(
												    shape=(self.vocab_size, self.d_embed)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_layers_._{i}_._weight"      ,  )
												__UpperCAmelCase	: str		=					self.add_weight(
												    shape=(self.vocab_size,)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_layers_._{i}_._bias"      ,  )
												self.out_layers.append((weight, bias))
						else:
									for i in range(len(self.cutoffs)):
												__UpperCAmelCase	: Tuple		=					self.cutoff_ends[i], self.cutoff_ends[i + 1]
												__UpperCAmelCase	: Any		=					self.d_embed // (self.div_val**i)

												__UpperCAmelCase	: Tuple		=					self.add_weight(
												    shape=(d_emb_i, self.d_proj)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_projs_._{i}")
												self.out_projs.append(UpperCamelCase_)
												__UpperCAmelCase	: Any		=					self.add_weight(
												    shape=(r_idx - l_idx, d_emb_i)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_layers_._{i}_._weight"      ,  )
												__UpperCAmelCase	: Any		=					self.add_weight(
												    shape=(r_idx - l_idx,)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_layers_._{i}_._bias"      ,  )
												self.out_layers.append((weight, bias))
						super().build(UpperCamelCase_)


			@staticmethod
			def       a_     (      UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Dict=None):


						"""simple docstring"""




						__UpperCAmelCase	: str		=					x
						if proj is not None:
									__UpperCAmelCase	: Optional[Any]		=					tf.einsum("ibd,ed->ibe"      ,  UpperCamelCase_      ,  UpperCamelCase_)
						return tf.einsum("ibd,nd->ibn"      ,  UpperCamelCase_      ,  UpperCamelCase_) + b


			@staticmethod
			def       a_     (      UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					shape_list(UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					tf.range(lp_size[0]      ,  dtype=target.dtype)
						__UpperCAmelCase	: List[Any]		=					tf.stack([r, target]      ,  1)
						return tf.gather_nd(UpperCamelCase_      ,  UpperCamelCase_)


			def       a_     (      self      :    Optional[Any]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Union[str, Any]=True      ,  UpperCamelCase_      :    str=False):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					0
						if self.n_clusters == 0:
									__UpperCAmelCase	: List[str]		=					self._logit(UpperCamelCase_      ,  self.out_layers[0][0]      ,  self.out_layers[0][1]      ,  self.out_projs[0])
									if target is not None:
												__UpperCAmelCase	: Any		=					tf.nn.sparse_softmax_cross_entropy_with_logits(labels=UpperCamelCase_      ,  logits=UpperCamelCase_)
									__UpperCAmelCase	: Tuple		=					tf.nn.log_softmax(UpperCamelCase_      ,  axis=-1)
						else:
									__UpperCAmelCase	: int		=					shape_list(UpperCamelCase_)
									__UpperCAmelCase	: Dict		=					[]
									__UpperCAmelCase	: List[Any]		=					tf.zeros(hidden_sizes[:2])
									for i in range(len(self.cutoffs)):
												__UpperCAmelCase	: Union[str, Any]		=					self.cutoff_ends[i], self.cutoff_ends[i + 1]
												if target is not None:
															__UpperCAmelCase	: int		=					(target >= l_idx) & (target < r_idx)
															__UpperCAmelCase	: Dict		=					tf.where(UpperCamelCase_)
															__UpperCAmelCase	: Union[str, Any]		=					tf.boolean_mask(UpperCamelCase_      ,  UpperCamelCase_) - l_idx

												if self.div_val == 1:
															__UpperCAmelCase	: int		=					self.out_layers[0][0][l_idx:r_idx]
															__UpperCAmelCase	: List[Any]		=					self.out_layers[0][1][l_idx:r_idx]
												else:
															__UpperCAmelCase	: Dict		=					self.out_layers[i][0]
															__UpperCAmelCase	: Dict		=					self.out_layers[i][1]

												if i == 0:
															__UpperCAmelCase	: Dict		=					tf.concat([cur_W, self.cluster_weight]      ,  0)
															__UpperCAmelCase	: int		=					tf.concat([cur_b, self.cluster_bias]      ,  0)

															__UpperCAmelCase	: Any		=					self._logit(UpperCamelCase_      ,  UpperCamelCase_      ,  UpperCamelCase_      ,  self.out_projs[0])
															__UpperCAmelCase	: Optional[Any]		=					tf.nn.log_softmax(UpperCamelCase_)
															out.append(head_logprob[..., : self.cutoffs[0]])
															if target is not None:
																		__UpperCAmelCase	: Optional[int]		=					tf.boolean_mask(UpperCamelCase_      ,  UpperCamelCase_)
																		__UpperCAmelCase	: List[Any]		=					self._gather_logprob(UpperCamelCase_      ,  UpperCamelCase_)
												else:
															__UpperCAmelCase	: Optional[int]		=					self._logit(UpperCamelCase_      ,  UpperCamelCase_      ,  UpperCamelCase_      ,  self.out_projs[i])
															__UpperCAmelCase	: Optional[int]		=					tf.nn.log_softmax(UpperCamelCase_)
															__UpperCAmelCase	: str		=					self.cutoffs[0] + i - 1  # No probability for the head cluster
															__UpperCAmelCase	: Tuple		=					head_logprob[..., cluster_prob_idx, None] + tail_logprob
															out.append(UpperCamelCase_)
															if target is not None:
																		__UpperCAmelCase	: Dict		=					tf.boolean_mask(UpperCamelCase_      ,  UpperCamelCase_)
																		__UpperCAmelCase	: Optional[Any]		=					tf.boolean_mask(UpperCamelCase_      ,  UpperCamelCase_)
																		__UpperCAmelCase	: Optional[int]		=					self._gather_logprob(UpperCamelCase_      ,  UpperCamelCase_)
																		cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
												if target is not None:
															loss += tf.scatter_nd(UpperCamelCase_      ,  -cur_logprob      ,  shape_list(UpperCamelCase_))
									__UpperCAmelCase	: List[Any]		=					tf.concat(UpperCamelCase_      ,  axis=-1)

						if target is not None:
									if return_mean:
												__UpperCAmelCase	: List[Any]		=					tf.reduce_mean(UpperCamelCase_)
									# Add the training-time loss value to the layer using `self.add_loss()`.
									self.add_loss(UpperCamelCase_)

									# Log the loss as a metric (we could log arbitrary metrics,
									# including different metrics for training and inference.
									self.add_metric(UpperCamelCase_      ,  name=self.name      ,  aggregation="mean" if return_mean else "")

						return out








"""simple docstring"""





from manim import *



class      a__				(	__magic_name__   ):


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					Rectangle(height=0.5      ,  width=0.5)
						__UpperCAmelCase	: List[str]		=					Rectangle(height=0.25      ,  width=0.25)
						__UpperCAmelCase	: Dict		=					Rectangle(height=0.46      ,  width=0.46).set_stroke(width=0)

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

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

						__UpperCAmelCase	: Any		=					[mem.copy() for i in range(6)]
						__UpperCAmelCase	: int		=					VGroup(*UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0)

						__UpperCAmelCase	: Union[str, Any]		=					Text("Model"      ,  font_size=24)
						__UpperCAmelCase	: Optional[Any]		=					Group(UpperCamelCase_      ,  UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0.5      ,  aligned_edge=UpperCamelCase_)
						model.move_to([3, -1.0, 0])
						self.add(UpperCamelCase_)

						__UpperCAmelCase	: str		=					[]
						__UpperCAmelCase	: Tuple		=					[]
						__UpperCAmelCase	: Any		=					[]

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

									__UpperCAmelCase	: Optional[int]		=					Rectangle(height=0.46 / 4      ,  width=0.46 / 3).set_stroke(width=0.0).set_fill(UpperCamelCase_      ,  opacity=0.7)

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

						self.add(*UpperCamelCase_      ,  *UpperCamelCase_      ,  *UpperCamelCase_)

						__UpperCAmelCase	: Optional[int]		=					[mem.copy() for i in range(6)]
						__UpperCAmelCase	: Dict		=					VGroup(*UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0)

						__UpperCAmelCase	: Tuple		=					Text("Loaded Checkpoint"      ,  font_size=24)
						__UpperCAmelCase	: Any		=					Group(UpperCamelCase_      ,  UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0.5      ,  aligned_edge=UpperCamelCase_)
						checkpoint.move_to([3, 0.5, 0])

						self.add(UpperCamelCase_)

						__UpperCAmelCase	: Optional[int]		=					[]
						__UpperCAmelCase	: int		=					[]

						for i, rect in enumerate(UpperCamelCase_):
									__UpperCAmelCase	: Optional[Any]		=					fill.copy().set_fill(UpperCamelCase_      ,  opacity=0.7)
									target.move_to(UpperCamelCase_)
									ckpt_arr.append(UpperCamelCase_)

									__UpperCAmelCase	: Optional[Any]		=					target.copy()
									if i < 5:
												cpu_target.move_to(cpu_left_col_base[i + 1])
									else:
												cpu_target.move_to(cpu_right_col_base[i - 5])
									ckpt_cpu_arr.append(UpperCamelCase_)
						self.add(*UpperCamelCase_      ,  *UpperCamelCase_)

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

						__UpperCAmelCase	: int		=					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(UpperCamelCase_      ,  UpperCamelCase_)

						__UpperCAmelCase	: Optional[Any]		=					MarkupText(
						    F"<span fgcolor='{BLUE}'>●</span> Checkpoint"      ,  font_size=18      ,  )

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

						__UpperCAmelCase	: Dict		=					MarkupText(
						    F"Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device."      ,  font_size=24      ,  )
						step_a.move_to([2, 2, 0])

						__UpperCAmelCase	: List[Any]		=					[meta_mem.copy() for i in range(6)]
						__UpperCAmelCase	: Optional[int]		=					[meta_mem.copy() for i in range(6)]
						__UpperCAmelCase	: int		=					VGroup(*UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0)
						__UpperCAmelCase	: Union[str, Any]		=					VGroup(*UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0)
						__UpperCAmelCase	: Tuple		=					VGroup(UpperCamelCase_      ,  UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0)
						__UpperCAmelCase	: Optional[int]		=					Text("Disk"      ,  font_size=24)
						__UpperCAmelCase	: str		=					Group(UpperCamelCase_      ,  UpperCamelCase_).arrange(UpperCamelCase_      ,  buff=0.5      ,  aligned_edge=UpperCamelCase_)
						disk.move_to([-4.0, -1.25, 0])
						self.play(Write(UpperCamelCase_      ,  run_time=3)      ,  Write(UpperCamelCase_      ,  run_time=1)      ,  Create(UpperCamelCase_      ,  run_time=1))

						__UpperCAmelCase	: str		=					[]
						for i, rect in enumerate(UpperCamelCase_):
									__UpperCAmelCase	: Union[str, Any]		=					rect.copy()
									target.generate_target()
									target.target.move_to(disk_left_col_base[i]).scale(0.5)
									animations.append(MoveToTarget(UpperCamelCase_      ,  run_time=1.5))
						self.play(*UpperCamelCase_)

						self.play(FadeOut(UpperCamelCase_))

						__UpperCAmelCase	: Union[str, Any]		=					MarkupText(F"Then, the checkpoint is removed from memory\nthrough garbage collection."      ,  font_size=24)
						step_a.move_to([2, 2, 0])

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

						self.play(
						    FadeOut(UpperCamelCase_      ,  UpperCamelCase_      ,  *UpperCamelCase_      ,  *UpperCamelCase_)      ,  )

						self.wait()


