






from manim import *




class  UpperCamelCase_      (			snake_case_				):



	'''simple docstring'''




	def 							_UpperCamelCase			(       self					)       ->    Tuple:
								snake_case_				=							Rectangle(height=0.5       ,			width=0.5					)
								snake_case_				=							Rectangle(height=0.46       ,			width=0.46					).set_stroke(width=0					)

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

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

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

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

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

								snake_case_				=							[]
								snake_case_				=							[]

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

															snake_case_				=							Rectangle(height=0.46       ,			width=0.46					).set_stroke(width=0.0					).set_fill(a       ,			opacity=0.8					)
															cpu_target.move_to(cpu_left_col_base[i]					)
															model_cpu_arr.append(a					)

								self.add(*a       ,			*a					)

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

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

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

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

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

								snake_case_				=							Square(0.3					)
								input.set_fill(a       ,			opacity=1.0					)
								input.set_stroke(width=0.0					)
								input.next_to(model_base[0]       ,			a       ,			buff=0.5					)

								self.play(Write(a					)					)

								input.generate_target()
								input.target.next_to(model_arr[0]       ,			direction=a       ,			buff=0.02					)
								self.play(MoveToTarget(a					)					)

								self.play(FadeOut(a					)					)

								snake_case_				=							Arrow(start=a       ,			end=a       ,			color=a       ,			buff=0.5					)
								a.next_to(model_arr[0].get_left()       ,			a       ,			buff=0.2					)

								model_cpu_arr[0].generate_target()
								model_cpu_arr[0].target.move_to(gpu_rect[0]					)

								snake_case_				=							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(a       ,			run_time=3					)					)

								snake_case_				=							{'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.02}

								self.play(
								    Write(a					)       ,			Circumscribe(model_arr[0]       ,			color=a       ,			**a					)       ,			Circumscribe(model_cpu_arr[0]       ,			color=a       ,			**a					)       ,			Circumscribe(gpu_rect[0]       ,			color=a       ,			**a					)       ,			)
								self.play(MoveToTarget(model_cpu_arr[0]					)					)

								snake_case_				=							a.copy()
								for i in range(6					):
															a_c.next_to(model_arr[i].get_right() + 0.02       ,			a       ,			buff=0.2					)

															input.generate_target()
															input.target.move_to(model_arr[i].get_right() + 0.02					)

															snake_case_				=							AnimationGroup(
															    FadeOut(a       ,			run_time=0.5					)       ,			MoveToTarget(a       ,			run_time=0.5					)       ,			FadeIn(a       ,			run_time=0.5					)       ,			lag_ratio=0.2					)

															self.play(a					)

															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:
																													snake_case_				=							0.7

																						self.play(
																						    Circumscribe(model_arr[i]       ,			**a					)       ,			Circumscribe(cpu_left_col_base[i]       ,			**a					)       ,			Circumscribe(cpu_left_col_base[i + 1]       ,			color=a       ,			**a					)       ,			Circumscribe(gpu_rect[0]       ,			color=a       ,			**a					)       ,			Circumscribe(model_arr[i + 1]       ,			color=a       ,			**a					)       ,			)
																						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=a       ,			**a					)       ,			Circumscribe(cpu_left_col_base[-1]       ,			color=a       ,			**a					)       ,			Circumscribe(gpu_rect[0]       ,			color=a       ,			**a					)       ,			)

																						self.play(MoveToTarget(model_cpu_arr[i]					)					)

															snake_case_				=							a_c
															snake_case_				=							a_c.copy()

								input.generate_target()
								input.target.next_to(model_base[-1]       ,			RIGHT + 0.02       ,			buff=0.5					)
								self.play(
								    FadeOut(a					)       ,			FadeOut(a       ,			run_time=0.5					)       ,			)

								snake_case_				=							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(a       ,			run_time=3					)       ,			MoveToTarget(a					)					)

								self.wait()













import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
	from transformers.models.roberta.modeling_flax_roberta import (
	    FlaxRobertaForCausalLM,
	    FlaxRobertaForMaskedLM,
	    FlaxRobertaForMultipleChoice,
	    FlaxRobertaForQuestionAnswering,
	    FlaxRobertaForSequenceClassification,
	    FlaxRobertaForTokenClassification,
	    FlaxRobertaModel,
	)




class  UpperCamelCase_      (			unittest.TestCase				):



	'''simple docstring'''




	def __init__(       self       ,			a       ,			a=13       ,			a=7       ,			a=True       ,			a=True       ,			a=True       ,			a=True       ,			a=99       ,			a=32       ,			a=5       ,			a=4       ,			a=37       ,			a="gelu"       ,			a=0.1       ,			a=0.1       ,			a=5_12       ,			a=16       ,			a=2       ,			a=0.02       ,			a=4       ,			)       ->    int:
								snake_case_				=							parent
								snake_case_				=							batch_size
								snake_case_				=							seq_length
								snake_case_				=							is_training
								snake_case_				=							use_attention_mask
								snake_case_				=							use_token_type_ids
								snake_case_				=							use_labels
								snake_case_				=							vocab_size
								snake_case_				=							hidden_size
								snake_case_				=							num_hidden_layers
								snake_case_				=							num_attention_heads
								snake_case_				=							intermediate_size
								snake_case_				=							hidden_act
								snake_case_				=							hidden_dropout_prob
								snake_case_				=							attention_probs_dropout_prob
								snake_case_				=							max_position_embeddings
								snake_case_				=							type_vocab_size
								snake_case_				=							type_sequence_label_size
								snake_case_				=							initializer_range
								snake_case_				=							num_choices




	def 							_UpperCamelCase			(       self					)       ->    List[str]:
								snake_case_				=							ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size					)

								snake_case_				=							None
								if self.use_attention_mask:
															snake_case_				=							random_attention_mask([self.batch_size, self.seq_length]					)

								snake_case_				=							None
								if self.use_token_type_ids:
															snake_case_				=							ids_tensor([self.batch_size, self.seq_length]       ,			self.type_vocab_size					)

								snake_case_				=							RobertaConfig(
								    vocab_size=self.vocab_size       ,			hidden_size=self.hidden_size       ,			num_hidden_layers=self.num_hidden_layers       ,			num_attention_heads=self.num_attention_heads       ,			intermediate_size=self.intermediate_size       ,			hidden_act=self.hidden_act       ,			hidden_dropout_prob=self.hidden_dropout_prob       ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,			max_position_embeddings=self.max_position_embeddings       ,			type_vocab_size=self.type_vocab_size       ,			is_decoder=a       ,			initializer_range=self.initializer_range       ,			)

								return config, input_ids, token_type_ids, attention_mask




	def 							_UpperCamelCase			(       self					)       ->    List[Any]:
								snake_case_				=							self.prepare_config_and_inputs()
								snake_case_			,		snake_case_			,		snake_case_			,		snake_case_				=							config_and_inputs
								snake_case_				=							{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
								return config, inputs_dict






	def 							_UpperCamelCase			(       self					)       ->    List[str]:
								snake_case_				=							self.prepare_config_and_inputs()
								snake_case_			,		snake_case_			,		snake_case_			,		snake_case_				=							config_and_inputs

								snake_case_				=							True
								snake_case_				=							floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
								snake_case_				=							ids_tensor([self.batch_size, self.seq_length]       ,			vocab_size=2					)

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)






@require_flax
class  UpperCamelCase_      (			snake_case_       ,					unittest.TestCase				):



	'''simple docstring'''




	lowerCAmelCase													= True

	lowerCAmelCase													= (
	    (
	        FlaxRobertaModel,
	        FlaxRobertaForCausalLM,
	        FlaxRobertaForMaskedLM,
	        FlaxRobertaForSequenceClassification,
	        FlaxRobertaForTokenClassification,
	        FlaxRobertaForMultipleChoice,
	        FlaxRobertaForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)




	def 							_UpperCamelCase			(       self					)       ->    List[str]:
								snake_case_				=							FlaxRobertaModelTester(self					)






	@slow
	def 							_UpperCamelCase			(       self					)       ->    str:
								for model_class_name in self.all_model_classes:
															snake_case_				=							model_class_name.from_pretrained('roberta-base'       ,			from_pt=a					)
															snake_case_				=							model(np.ones((1, 1)					)					)
															self.assertIsNotNone(a					)





