


'''simple docstring'''
import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import DecisionTransformerModel
  from transformers.models.decision_transformer.modeling_decision_transformer import (
      DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )




class      lowerCamelCase_  :






      def __init__(   self							:       List[str]		,		_A							:       List[Any]		,		_A							:       int=13		,		_A							:       List[Any]=7		,		_A							:       Optional[Any]=6		,		_A							:       Tuple=17		,		_A							:       Union[str, Any]=23		,		_A							:       str=11		,		_A							:       Dict=True		,		):


          '''simple docstring'''


          UpperCAmelCase__     :						Optional[int]						=							parent
          UpperCAmelCase__     :						Tuple						=							batch_size
          UpperCAmelCase__     :						Optional[Any]						=							seq_length
          UpperCAmelCase__     :						Tuple						=							act_dim
          UpperCAmelCase__     :						Dict						=							state_dim
          UpperCAmelCase__     :						int						=							hidden_size
          UpperCAmelCase__     :						Dict						=							max_length
          UpperCAmelCase__     :						Any						=							is_training






      def lowercase_				(   self							:       Dict					):


          '''simple docstring'''


          UpperCAmelCase__     :						List[str]						=							floats_tensor((self.batch_size, self.seq_length, self.state_dim)					)
          UpperCAmelCase__     :						Dict						=							floats_tensor((self.batch_size, self.seq_length, self.act_dim)					)
          UpperCAmelCase__     :						Optional[Any]						=							floats_tensor((self.batch_size, self.seq_length, 1)					)
          UpperCAmelCase__     :						str						=							floats_tensor((self.batch_size, self.seq_length, 1)					)
          UpperCAmelCase__     :						str						=							ids_tensor((self.batch_size, self.seq_length)		,		vocab_size=1_000					)
          UpperCAmelCase__     :						Optional[Any]						=							random_attention_mask((self.batch_size, self.seq_length)					)

          UpperCAmelCase__     :						Tuple						=							self.get_config()

          return (
              config,
              states,
              actions,
              rewards,
              returns_to_go,
              timesteps,
              attention_mask,
          )






      def lowercase_				(   self							:       Dict					):


          '''simple docstring'''


          return DecisionTransformerConfig(
              batch_size=self.batch_size		,		seq_length=self.seq_length		,		act_dim=self.act_dim		,		state_dim=self.state_dim		,		hidden_size=self.hidden_size		,		max_length=self.max_length		,		)






      def lowercase_				(   self							:       int		,		_A							:       int		,		_A							:       Optional[int]		,		_A							:       int		,		_A							:       int		,		_A							:       Union[str, Any]		,		_A							:       List[str]		,		_A							:       List[str]		,		):


          '''simple docstring'''


          UpperCAmelCase__     :						Union[str, Any]						=							DecisionTransformerModel(config=_A					)
          model.to(_A					)
          model.eval()
          UpperCAmelCase__     :						List[Any]						=							model(_A		,		_A		,		_A		,		_A		,		_A		,		_A					)

          self.parent.assertEqual(result.state_preds.shape		,		states.shape					)
          self.parent.assertEqual(result.action_preds.shape		,		actions.shape					)
          self.parent.assertEqual(result.return_preds.shape		,		returns_to_go.shape					)
          self.parent.assertEqual(
              result.last_hidden_state.shape		,		(self.batch_size, self.seq_length * 3, self.hidden_size)					)  # seq length *3 as there are 3 modelities: states, returns and actions






      def lowercase_				(   self							:       Any					):


          '''simple docstring'''


          UpperCAmelCase__     :						Tuple						=							self.prepare_config_and_inputs()
          (
              (
              UpperCAmelCase__
          )		,		(
              UpperCAmelCase__
          )		,		(
              UpperCAmelCase__
          )		,		(
              UpperCAmelCase__
          )		,		(
              UpperCAmelCase__
          )		,		(
              UpperCAmelCase__
          )		,		(
              UpperCAmelCase__
          )		,		
          )     :						Union[str, Any]						=							config_and_inputs
          UpperCAmelCase__     :						Optional[int]						=							{
              '''states''': states,
              '''actions''': actions,
              '''rewards''': rewards,
              '''returns_to_go''': returns_to_go,
              '''timesteps''': timesteps,
              '''attention_mask''': attention_mask,
          }
          return config, inputs_dict




@require_torch
class      lowerCamelCase_  (      __a	,			__a	,			__a	,			unittest.TestCase ):
      lowerCAmelCase__							=	(DecisionTransformerModel,) if is_torch_available() else ()
      lowerCAmelCase__							=	()
      lowerCAmelCase__							=	{'feature-extraction': DecisionTransformerModel} if is_torch_available() else {}

      # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
      lowerCAmelCase__							=	False

      # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False
      lowerCAmelCase__							=	False






      def lowercase_				(   self							:       Optional[Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						List[Any]						=							DecisionTransformerModelTester(self					)
          UpperCAmelCase__     :						str						=							ConfigTester(self		,		config_class=_A		,		hidden_size=37					)






      def lowercase_				(   self							:       Optional[int]					):


          '''simple docstring'''


          self.config_tester.run_common_tests()






      def lowercase_				(   self							:       List[Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						Tuple						=							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_A					)






      @slow
      def lowercase_				(   self							:       Tuple					):


          '''simple docstring'''


          for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase__     :						List[str]						=							DecisionTransformerModel.from_pretrained(_A					)
              self.assertIsNotNone(_A					)






      def lowercase_				(   self							:       Optional[Any]					):


          '''simple docstring'''


          UpperCAmelCase__		,		UpperCAmelCase__     :						List[Any]						=							self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              UpperCAmelCase__     :						int						=							model_class(_A					)
              UpperCAmelCase__     :						Tuple						=							inspect.signature(model.forward					)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              UpperCAmelCase__     :						Dict						=							[*signature.parameters.keys()]

              UpperCAmelCase__     :						str						=							[
                  '''states''',
                  '''actions''',
                  '''rewards''',
                  '''returns_to_go''',
                  '''timesteps''',
                  '''attention_mask''',
              ]

              self.assertListEqual(arg_names[: len(_A					)]		,		_A					)




@require_torch
class      lowerCamelCase_  (      unittest.TestCase ):






      @slow
      def lowercase_				(   self							:       List[Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						Dict						=							2  # number of steps of autoregressive prediction we will perform
          UpperCAmelCase__     :						Dict						=							10  # defined by the RL environment, may be normalized
          UpperCAmelCase__     :						Optional[int]						=							DecisionTransformerModel.from_pretrained('''edbeeching/decision-transformer-gym-hopper-expert'''					)
          UpperCAmelCase__     :						Union[str, Any]						=							model.to(_A					)
          UpperCAmelCase__     :						str						=							model.config
          torch.manual_seed(0					)
          UpperCAmelCase__     :						str						=							torch.randn(1		,		1		,		config.state_dim					).to(device=_A		,		dtype=torch.floataa					)  # env.reset()

          UpperCAmelCase__     :						Optional[int]						=							torch.tensor(
              [[0.2_4_2_7_9_3, -0.2_8_6_9_3_0_7_4, 0.8_7_4_2_6_1_3], [0.6_7_8_1_5_2_7_4, -0.0_8_1_0_1_0_8_5, -0.1_2_9_5_2_1_4_7]]		,		device=_A					)

          UpperCAmelCase__     :						Tuple						=							torch.tensor(_A		,		device=_A		,		dtype=torch.floataa					).reshape(1		,		1		,		1					)
          UpperCAmelCase__     :						Tuple						=							state
          UpperCAmelCase__     :						Union[str, Any]						=							torch.zeros(1		,		0		,		config.act_dim		,		device=_A		,		dtype=torch.floataa					)
          UpperCAmelCase__     :						int						=							torch.zeros(1		,		0		,		device=_A		,		dtype=torch.floataa					)
          UpperCAmelCase__     :						Union[str, Any]						=							torch.tensor(0		,		device=_A		,		dtype=torch.long					).reshape(1		,		1					)

          for step in range(_A					):
              UpperCAmelCase__     :						Optional[Any]						=							torch.cat([actions, torch.zeros(1		,		1		,		config.act_dim		,		device=_A					)]		,		dim=1					)
              UpperCAmelCase__     :						Union[str, Any]						=							torch.cat([rewards, torch.zeros(1		,		1		,		device=_A					)]		,		dim=1					)

              UpperCAmelCase__     :						List[str]						=							torch.ones(1		,		states.shape[1]					).to(dtype=torch.long		,		device=states.device					)

              with torch.no_grad():
                  UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__     :						Any						=							model(
                      states=_A		,		actions=_A		,		rewards=_A		,		returns_to_go=_A		,		timesteps=_A		,		attention_mask=_A		,		return_dict=_A		,		)

              self.assertEqual(action_pred.shape		,		actions.shape					)
              self.assertTrue(torch.allclose(action_pred[0, -1]		,		expected_outputs[step]		,		atol=1e-4					)					)
              UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__     :						Dict						=							(  # env.step(action)
                  torch.randn(1		,		1		,		config.state_dim					).to(device=_A		,		dtype=torch.floataa					),
                  1.0,
                  False,
                  {},
              )

              UpperCAmelCase__     :						Optional[Any]						=							action_pred[0, -1]
              UpperCAmelCase__     :						Tuple						=							torch.cat([states, state]		,		dim=1					)
              UpperCAmelCase__     :						Any						=							returns_to_go[0, -1] - reward
              UpperCAmelCase__     :						Union[str, Any]						=							torch.cat([returns_to_go, pred_return.reshape(1		,		1		,		1					)]		,		dim=1					)
              UpperCAmelCase__     :						Optional[int]						=							torch.cat(
                  [timesteps, torch.ones((1, 1)		,		device=_A		,		dtype=torch.long					) * (step + 1)]		,		dim=1					)






"""simple docstring"""

from manim import *





class 							lowercase_			(	__lowerCAmelCase				):



 '''simple docstring'''

 def      lowerCAmelCase_				(					self			:			Union[str, Any]						):
       _A     =  Rectangle(height=0.5       ,				width=0.5						)
       _A     =  Rectangle(height=0.46       ,				width=0.46						).set_stroke(width=0						)

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

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

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

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

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

       _A     =  []
       _A     =  []

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

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

       self.add(*_UpperCAmelCase       ,				*_UpperCAmelCase						)

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

       _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(_UpperCAmelCase       ,				_UpperCAmelCase						)

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

       _A     =  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(_UpperCAmelCase						)						)

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

       self.play(Write(_UpperCAmelCase						)						)

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

       self.play(FadeOut(_UpperCAmelCase						)						)

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

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

       _A     =  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(_UpperCAmelCase       ,				run_time=3						)						)

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

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

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

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

             _A     =  AnimationGroup(
                 FadeOut(_UpperCAmelCase       ,				run_time=0.5						)       ,				MoveToTarget(_UpperCAmelCase       ,				run_time=0.5						)       ,				FadeIn(_UpperCAmelCase       ,				run_time=0.5						)       ,				lag_ratio=0.2						)

             self.play(_UpperCAmelCase						)

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

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

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

             _A     =  a_c
             _A     =  a_c.copy()

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

       _A     =  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(_UpperCAmelCase       ,				run_time=3						)       ,				MoveToTarget(_UpperCAmelCase						)						)

       self.wait()

