


import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
	import jax.numpy as jnp

if is_tf_available():
	import tensorflow as tf

if is_torch_available():
	import torch




class       lowercase_					(unittest.TestCase   ):





				def 						__UpperCamelCase    (				self)  -> Dict:
											a__      ={
											    'task_specific_params': {
											        'summarization': {'length_penalty': 1.0, 'max_length': 128, 'min_length': 12, 'num_beams': 4},
											        'summarization_cnn': {'length_penalty': 2.0, 'max_length': 142, 'min_length': 56, 'num_beams': 4},
											        'summarization_xsum': {'length_penalty': 1.0, 'max_length': 62, 'min_length': 11, 'num_beams': 6},
											    }
											}
											a__      ={
											    'task_specific_params.summarization.length_penalty': 1.0,
											    'task_specific_params.summarization.max_length': 128,
											    'task_specific_params.summarization.min_length': 12,
											    'task_specific_params.summarization.num_beams': 4,
											    'task_specific_params.summarization_cnn.length_penalty': 2.0,
											    'task_specific_params.summarization_cnn.max_length': 142,
											    'task_specific_params.summarization_cnn.min_length': 56,
											    'task_specific_params.summarization_cnn.num_beams': 4,
											    'task_specific_params.summarization_xsum.length_penalty': 1.0,
											    'task_specific_params.summarization_xsum.max_length': 62,
											    'task_specific_params.summarization_xsum.min_length': 11,
											    'task_specific_params.summarization_xsum.num_beams': 6,
											}

											self.assertEqual(flatten_dict(lowercase_)       ,  lowercase_)





				def 						__UpperCamelCase    (				self)  -> Any:
											a__      =np.random.randn(3       ,  4)
											self.assertTrue(np.allclose(transpose(lowercase_)       ,  x.transpose()))

											a__      =np.random.randn(3       ,  4       ,  5)
											self.assertTrue(np.allclose(transpose(lowercase_       ,  axes=(1, 2, 0))       ,  x.transpose((1, 2, 0))))





				@require_torch
				def 						__UpperCamelCase    (				self)  -> List[str]:
											a__      =np.random.randn(3       ,  4)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(transpose(lowercase_)       ,  transpose(lowercase_).numpy()))

											a__      =np.random.randn(3       ,  4       ,  5)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(transpose(lowercase_       ,  axes=(1, 2, 0))       ,  transpose(lowercase_       ,  axes=(1, 2, 0)).numpy()))





				@require_tf
				def 						__UpperCamelCase    (				self)  -> List[str]:
											a__      =np.random.randn(3       ,  4)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(transpose(lowercase_)       ,  transpose(lowercase_).numpy()))

											a__      =np.random.randn(3       ,  4       ,  5)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(transpose(lowercase_       ,  axes=(1, 2, 0))       ,  transpose(lowercase_       ,  axes=(1, 2, 0)).numpy()))





				@require_flax
				def 						__UpperCamelCase    (				self)  -> Optional[int]:
											a__      =np.random.randn(3       ,  4)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(transpose(lowercase_)       ,  np.asarray(transpose(lowercase_))))

											a__      =np.random.randn(3       ,  4       ,  5)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(transpose(lowercase_       ,  axes=(1, 2, 0))       ,  np.asarray(transpose(lowercase_       ,  axes=(1, 2, 0)))))





				def 						__UpperCamelCase    (				self)  -> List[str]:
											a__      =np.random.randn(3       ,  4)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (4, 3))       ,  np.reshape(lowercase_       ,  (4, 3))))

											a__      =np.random.randn(3       ,  4       ,  5)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (12, 5))       ,  np.reshape(lowercase_       ,  (12, 5))))





				@require_torch
				def 						__UpperCamelCase    (				self)  -> int:
											a__      =np.random.randn(3       ,  4)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (4, 3))       ,  reshape(lowercase_       ,  (4, 3)).numpy()))

											a__      =np.random.randn(3       ,  4       ,  5)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (12, 5))       ,  reshape(lowercase_       ,  (12, 5)).numpy()))





				@require_tf
				def 						__UpperCamelCase    (				self)  -> Dict:
											a__      =np.random.randn(3       ,  4)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (4, 3))       ,  reshape(lowercase_       ,  (4, 3)).numpy()))

											a__      =np.random.randn(3       ,  4       ,  5)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (12, 5))       ,  reshape(lowercase_       ,  (12, 5)).numpy()))





				@require_flax
				def 						__UpperCamelCase    (				self)  -> Any:
											a__      =np.random.randn(3       ,  4)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (4, 3))       ,  np.asarray(reshape(lowercase_       ,  (4, 3)))))

											a__      =np.random.randn(3       ,  4       ,  5)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(reshape(lowercase_       ,  (12, 5))       ,  np.asarray(reshape(lowercase_       ,  (12, 5)))))





				def 						__UpperCamelCase    (				self)  -> Tuple:
											a__      =np.random.randn(1       ,  3       ,  4)
											self.assertTrue(np.allclose(squeeze(lowercase_)       ,  np.squeeze(lowercase_)))

											a__      =np.random.randn(1       ,  4       ,  1       ,  5)
											self.assertTrue(np.allclose(squeeze(lowercase_       ,  axis=2)       ,  np.squeeze(lowercase_       ,  axis=2)))





				@require_torch
				def 						__UpperCamelCase    (				self)  -> Dict:
											a__      =np.random.randn(1       ,  3       ,  4)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(squeeze(lowercase_)       ,  squeeze(lowercase_).numpy()))

											a__      =np.random.randn(1       ,  4       ,  1       ,  5)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(squeeze(lowercase_       ,  axis=2)       ,  squeeze(lowercase_       ,  axis=2).numpy()))





				@require_tf
				def 						__UpperCamelCase    (				self)  -> List[Any]:
											a__      =np.random.randn(1       ,  3       ,  4)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(squeeze(lowercase_)       ,  squeeze(lowercase_).numpy()))

											a__      =np.random.randn(1       ,  4       ,  1       ,  5)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(squeeze(lowercase_       ,  axis=2)       ,  squeeze(lowercase_       ,  axis=2).numpy()))





				@require_flax
				def 						__UpperCamelCase    (				self)  -> int:
											a__      =np.random.randn(1       ,  3       ,  4)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(squeeze(lowercase_)       ,  np.asarray(squeeze(lowercase_))))

											a__      =np.random.randn(1       ,  4       ,  1       ,  5)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(squeeze(lowercase_       ,  axis=2)       ,  np.asarray(squeeze(lowercase_       ,  axis=2))))





				def 						__UpperCamelCase    (				self)  -> Optional[Any]:
											a__      =np.random.randn(3       ,  4)
											self.assertTrue(np.allclose(expand_dims(lowercase_       ,  axis=1)       ,  np.expand_dims(lowercase_       ,  axis=1)))





				@require_torch
				def 						__UpperCamelCase    (				self)  -> Union[str, Any]:
											a__      =np.random.randn(3       ,  4)
											a__      =torch.tensor(lowercase_)
											self.assertTrue(np.allclose(expand_dims(lowercase_       ,  axis=1)       ,  expand_dims(lowercase_       ,  axis=1).numpy()))





				@require_tf
				def 						__UpperCamelCase    (				self)  -> Tuple:
											a__      =np.random.randn(3       ,  4)
											a__      =tf.constant(lowercase_)
											self.assertTrue(np.allclose(expand_dims(lowercase_       ,  axis=1)       ,  expand_dims(lowercase_       ,  axis=1).numpy()))





				@require_flax
				def 						__UpperCamelCase    (				self)  -> Optional[Any]:
											a__      =np.random.randn(3       ,  4)
											a__      =jnp.array(lowercase_)
											self.assertTrue(np.allclose(expand_dims(lowercase_       ,  axis=1)       ,  np.asarray(expand_dims(lowercase_       ,  axis=1))))





from manim import *




class       lowercase_					(lowercase__   ):





				def 						__UpperCamelCase    (				self)  -> List[Any]:
											a__      =Rectangle(height=0.5       ,  width=0.5)
											a__      =Rectangle(height=0.46       ,  width=0.46).set_stroke(width=0)

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

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

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

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

											a__      =[]
											for i, rect in enumerate(lowercase_):
																		rect.set_stroke(lowercase_)
																		# target = fill.copy().set_fill(YELLOW, opacity=0.7)
																		# target.move_to(rect)
																		# self.add(target)

																		a__      =Rectangle(height=0.46 / 4       ,  width=0.46 / 3).set_stroke(width=0.0).set_fill(lowercase_       ,  opacity=0.7)

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

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

											a__      =Text('Loaded Checkpoint'       ,  font_size=24)
											a__      =Group(lowercase_       ,  lowercase_).arrange(lowercase_       ,  aligned_edge=lowercase_       ,  buff=0.4)
											checkpoint.move_to([3, 0.5, 0])

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

											a__      =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(lowercase_       ,  lowercase_)

											a__      =MarkupText(
											    F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""       ,  font_size=18       ,  )

											blue_text.next_to(lowercase_       ,  DOWN * 2.4       ,  aligned_edge=key_text.get_left())

											a__      =MarkupText(
											    F"""Next, a <i><span fgcolor=\"{BLUE}\">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor=\"{BLUE}\">single shard</span>."""       ,  font_size=24       ,  )
											step_a.move_to([2, 2, 0])
											self.play(Write(lowercase_)       ,  Write(lowercase_))

											self.play(Write(lowercase_       ,  run_time=1)       ,  Create(lowercase_       ,  run_time=1))

											a__      =[]
											a__      =[]
											for i, rect in enumerate(lowercase_):
																		a__      =fill.copy().set_fill(lowercase_       ,  opacity=0.7)
																		target.move_to(lowercase_)
																		first_animations.append(GrowFromCenter(lowercase_       ,  run_time=1))

																		a__      =target.copy()
																		cpu_target.generate_target()
																		if i < 5:
																									cpu_target.target.move_to(cpu_left_col_base[i + 1])
																		else:
																									cpu_target.target.move_to(cpu_right_col_base[i - 5])
																		second_animations.append(MoveToTarget(lowercase_       ,  run_time=1.5))

											self.play(*lowercase_)
											self.play(*lowercase_)
											self.wait()

