






'''simple docstring'''

from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCamelCase      					=logging.get_logger(__name__)

# General docstring
UpperCamelCase      					="RegNetConfig"

# Base docstring
UpperCamelCase      					="facebook/regnet-y-040"
UpperCamelCase      					=[1, 1088, 7, 7]

# Image classification docstring
UpperCamelCase      					="facebook/regnet-y-040"
UpperCamelCase      					="tabby, tabby cat"

UpperCamelCase      					=[
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]




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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = 3      ,  __lowerCAmelCase = 1      ,  __lowerCAmelCase = 1      ,  __lowerCAmelCase = "relu"      ,  **__lowerCAmelCase      ,  ):
       super().__init__(**__lowerCAmelCase			)
       # The padding and conv has been verified in
       # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
       UpperCamelCase_					: Any            =				tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2			)
       UpperCamelCase_					: Optional[int]            =				tf.keras.layers.ConvaD(
           filters=__lowerCAmelCase      ,  kernel_size=__lowerCAmelCase      ,  strides=__lowerCAmelCase      ,  padding="""VALID"""      ,  groups=__lowerCAmelCase      ,  use_bias=__lowerCAmelCase      ,  name="""convolution"""      ,  )
       UpperCamelCase_					: Any            =				tf.keras.layers.BatchNormalization(epsilon=1E-5      ,  momentum=0.9      ,  name="""normalization"""			)
       UpperCamelCase_					: List[Any]            =				ACTaFN[activation] if activation is not None else tf.identity





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase			):
       UpperCamelCase_					: List[Any]            =				self.convolution(self.padding(__lowerCAmelCase			)			)
       UpperCamelCase_					: int            =				self.normalization(__lowerCAmelCase			)
       UpperCamelCase_					: List[Any]            =				self.activation(__lowerCAmelCase			)
       return hidden_state






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: Dict            =				config.num_channels
       UpperCamelCase_					: Any            =				TFRegNetConvLayer(
           out_channels=config.embedding_size      ,  kernel_size=3      ,  stride=2      ,  activation=config.hidden_act      ,  name="""embedder"""      ,  )





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase			):
       UpperCamelCase_					: Optional[Any]            =				shape_list(__lowerCAmelCase			)[1]
       if tf.executing_eagerly() and num_channels != self.num_channels:
             raise ValueError(
                 """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""			)

       # When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
       # So change the input format from `NCHW` to `NHWC`.
       # shape = (batch_size, in_height, in_width, in_channels=num_channels)
       UpperCamelCase_					: List[str]            =				tf.transpose(__lowerCAmelCase      ,  perm=(0, 2, 3, 1)			)
       UpperCamelCase_					: List[str]            =				self.embedder(__lowerCAmelCase			)
       return hidden_state






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = 2      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: List[Any]            =				tf.keras.layers.ConvaD(
           filters=__lowerCAmelCase      ,  kernel_size=1      ,  strides=__lowerCAmelCase      ,  use_bias=__lowerCAmelCase      ,  name="""convolution"""			)
       UpperCamelCase_					: int            =				tf.keras.layers.BatchNormalization(epsilon=1E-5      ,  momentum=0.9      ,  name="""normalization"""			)





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = False			):
       return self.normalization(self.convolution(__lowerCAmelCase			)      ,  training=__lowerCAmelCase			)






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: List[Any]            =				tf.keras.layers.GlobalAveragePoolingaD(keepdims=__lowerCAmelCase      ,  name="""pooler"""			)
       UpperCamelCase_					: List[str]            =				[
           tf.keras.layers.ConvaD(filters=__lowerCAmelCase      ,  kernel_size=1      ,  activation="""relu"""      ,  name="""attention.0"""			),
           tf.keras.layers.ConvaD(filters=__lowerCAmelCase      ,  kernel_size=1      ,  activation="""sigmoid"""      ,  name="""attention.2"""			),
       ]





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase			):
       # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
       UpperCamelCase_					: Optional[int]            =				self.pooler(__lowerCAmelCase			)
       for layer_module in self.attention:
             UpperCamelCase_					: Dict            =				layer_module(__lowerCAmelCase			)
       UpperCamelCase_					: Any            =				hidden_state * pooled
       return hidden_state






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase = 1      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: Union[str, Any]            =				in_channels != out_channels or stride != 1
       UpperCamelCase_					: List[Any]            =				max(1      ,  out_channels // config.groups_width			)
       UpperCamelCase_					: List[Any]            =				(
           TFRegNetShortCut(__lowerCAmelCase      ,  stride=__lowerCAmelCase      ,  name="""shortcut"""			)
           if should_apply_shortcut
           else tf.keras.layers.Activation("""linear"""      ,  name="""shortcut"""			)
       )
       # `self.layers` instead of `self.layer` because that is a reserved argument.
       UpperCamelCase_					: str            =				[
           TFRegNetConvLayer(__lowerCAmelCase      ,  kernel_size=1      ,  activation=config.hidden_act      ,  name="""layer.0"""			),
           TFRegNetConvLayer(
               __lowerCAmelCase      ,  stride=__lowerCAmelCase      ,  groups=__lowerCAmelCase      ,  activation=config.hidden_act      ,  name="""layer.1"""			),
           TFRegNetConvLayer(__lowerCAmelCase      ,  kernel_size=1      ,  activation=__lowerCAmelCase      ,  name="""layer.2"""			),
       ]
       UpperCamelCase_					: List[str]            =				ACTaFN[config.hidden_act]





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase			):
       UpperCamelCase_					: Optional[Any]            =				hidden_state
       for layer_module in self.layers:
             UpperCamelCase_					: Union[str, Any]            =				layer_module(__lowerCAmelCase			)
       UpperCamelCase_					: List[Any]            =				self.shortcut(__lowerCAmelCase			)
       hidden_state += residual
       UpperCamelCase_					: List[str]            =				self.activation(__lowerCAmelCase			)
       return hidden_state






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase = 1      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: Dict            =				in_channels != out_channels or stride != 1
       UpperCamelCase_					: Optional[Any]            =				max(1      ,  out_channels // config.groups_width			)
       UpperCamelCase_					: Tuple            =				(
           TFRegNetShortCut(__lowerCAmelCase      ,  stride=__lowerCAmelCase      ,  name="""shortcut"""			)
           if should_apply_shortcut
           else tf.keras.layers.Activation("""linear"""      ,  name="""shortcut"""			)
       )
       UpperCamelCase_					: List[str]            =				[
           TFRegNetConvLayer(__lowerCAmelCase      ,  kernel_size=1      ,  activation=config.hidden_act      ,  name="""layer.0"""			),
           TFRegNetConvLayer(
               __lowerCAmelCase      ,  stride=__lowerCAmelCase      ,  groups=__lowerCAmelCase      ,  activation=config.hidden_act      ,  name="""layer.1"""			),
           TFRegNetSELayer(__lowerCAmelCase      ,  reduced_channels=int(round(in_channels / 4			)			)      ,  name="""layer.2"""			),
           TFRegNetConvLayer(__lowerCAmelCase      ,  kernel_size=1      ,  activation=__lowerCAmelCase      ,  name="""layer.3"""			),
       ]
       UpperCamelCase_					: Dict            =				ACTaFN[config.hidden_act]





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase			):
       UpperCamelCase_					: Dict            =				hidden_state
       for layer_module in self.layers:
             UpperCamelCase_					: Union[str, Any]            =				layer_module(__lowerCAmelCase			)
       UpperCamelCase_					: Any            =				self.shortcut(__lowerCAmelCase			)
       hidden_state += residual
       UpperCamelCase_					: Optional[int]            =				self.activation(__lowerCAmelCase			)
       return hidden_state






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase = 2      ,  __lowerCAmelCase = 2      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)

       UpperCamelCase_					: Optional[int]            =				TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
       UpperCamelCase_					: Union[str, Any]            =				[
           # downsampling is done in the first layer with stride of 2
           layer(__lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  stride=__lowerCAmelCase      ,  name="""layers.0"""			),
           *[layer(__lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  name=F"layers.{i+1}"			) for i in range(depth - 1			)],
       ]





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase			):
       for layer_module in self.layers:
             UpperCamelCase_					: List[Any]            =				layer_module(__lowerCAmelCase			)
       return hidden_state






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






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: Tuple            =				[]
       # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
       self.stages.append(
           TFRegNetStage(
               __lowerCAmelCase      ,  config.embedding_size      ,  config.hidden_sizes[0]      ,  stride=2 if config.downsample_in_first_stage else 1      ,  depth=config.depths[0]      ,  name="""stages.0"""      ,  )			)
       UpperCamelCase_					: Optional[int]            =				zip(config.hidden_sizes      ,  config.hidden_sizes[1:]			)
       for i, ((in_channels, out_channels), depth) in enumerate(zip(__lowerCAmelCase      ,  config.depths[1:]			)			):
             self.stages.append(TFRegNetStage(__lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  depth=__lowerCAmelCase      ,  name=F"stages.{i+1}"			)			)





 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = False      ,  __lowerCAmelCase = True			):
       UpperCamelCase_					: Optional[int]            =				() if output_hidden_states else None

       for stage_module in self.stages:
             if output_hidden_states:
                   UpperCamelCase_					: List[str]            =				hidden_states + (hidden_state,)

             UpperCamelCase_					: Optional[Any]            =				stage_module(__lowerCAmelCase			)

       if output_hidden_states:
             UpperCamelCase_					: Tuple            =				hidden_states + (hidden_state,)

       if not return_dict:
             return tuple(v for v in [hidden_state, hidden_states] if v is not None			)

       return TFBaseModelOutputWithNoAttention(last_hidden_state=__lowerCAmelCase      ,  hidden_states=__lowerCAmelCase			)






@keras_serializable
class      A     (   tf.keras.layers.Layer    ):






 """simple docstring"""





 __a      :    Optional[Any]			  =    RegNetConfig
 def __init__(				self      ,  __lowerCAmelCase      ,  **__lowerCAmelCase			):
       super().__init__(**__lowerCAmelCase			)
       UpperCamelCase_					: Union[str, Any]            =				config
       UpperCamelCase_					: Dict            =				TFRegNetEmbeddings(__lowerCAmelCase      ,  name="""embedder"""			)
       UpperCamelCase_					: str            =				TFRegNetEncoder(__lowerCAmelCase      ,  name="""encoder"""			)
       UpperCamelCase_					: List[Any]            =				tf.keras.layers.GlobalAveragePoolingaD(keepdims=__lowerCAmelCase      ,  name="""pooler"""			)





 @unpack_inputs
 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = False      ,  ):
       UpperCamelCase_					: Dict            =				(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       UpperCamelCase_					: str            =				return_dict if return_dict is not None else self.config.use_return_dict

       UpperCamelCase_					: Union[str, Any]            =				self.embedder(__lowerCAmelCase      ,  training=__lowerCAmelCase			)

       UpperCamelCase_					: int            =				self.encoder(
           __lowerCAmelCase      ,  output_hidden_states=__lowerCAmelCase      ,  return_dict=__lowerCAmelCase      ,  training=__lowerCAmelCase			)

       UpperCamelCase_					: Any            =				encoder_outputs[0]
       UpperCamelCase_					: List[Any]            =				self.pooler(__lowerCAmelCase			)

       # Change to NCHW output format have uniformity in the modules
       UpperCamelCase_					: Tuple            =				tf.transpose(__lowerCAmelCase      ,  perm=(0, 3, 1, 2)			)
       UpperCamelCase_					: Optional[Any]            =				tf.transpose(__lowerCAmelCase      ,  perm=(0, 3, 1, 2)			)

       # Change the other hidden state outputs to NCHW as well
       if output_hidden_states:
             UpperCamelCase_					: List[Any]            =				tuple([tf.transpose(__lowerCAmelCase      ,  perm=(0, 3, 1, 2)			) for h in encoder_outputs[1]]			)

       if not return_dict:
             return (last_hidden_state, pooled_output) + encoder_outputs[1:]

       return TFBaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=__lowerCAmelCase      ,  pooler_output=__lowerCAmelCase      ,  hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states      ,  )






class      A     (   SCREAMING_SNAKE_CASE__    ):






 """simple docstring"""





 __a      :    int			  =    RegNetConfig
 __a      :    Tuple			  =    '''regnet'''
 __a      :    Optional[int]			  =    '''pixel_values'''
 @property
 def 	_UpperCAmelCase							(				self			):
       return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)      ,  dtype=tf.floataa			)}






UpperCamelCase      					=R"\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n"

UpperCamelCase      					=R"\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"




@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.''', SCREAMING_SNAKE_CASE__, )
class      A     (   SCREAMING_SNAKE_CASE__    ):






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  *__lowerCAmelCase      ,  **__lowerCAmelCase			):
       super().__init__(__lowerCAmelCase      ,  *__lowerCAmelCase      ,  **__lowerCAmelCase			)
       UpperCamelCase_					: str            =				TFRegNetMainLayer(__lowerCAmelCase      ,  name="""regnet"""			)





 @unpack_inputs
 @add_start_docstrings_to_model_forward(__lowerCAmelCase			)
 @add_code_sample_docstrings(
     checkpoint=_CHECKPOINT_FOR_DOC      ,  output_type=__lowerCAmelCase      ,  config_class=_CONFIG_FOR_DOC      ,  modality="""vision"""      ,  expected_output=_EXPECTED_OUTPUT_SHAPE      ,  )
 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase=False      ,  ):
       UpperCamelCase_					: Optional[int]            =				(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       UpperCamelCase_					: List[Any]            =				return_dict if return_dict is not None else self.config.use_return_dict

       UpperCamelCase_					: Union[str, Any]            =				self.regnet(
           pixel_values=__lowerCAmelCase      ,  output_hidden_states=__lowerCAmelCase      ,  return_dict=__lowerCAmelCase      ,  training=__lowerCAmelCase      ,  )
       if not return_dict:
             return (outputs[0],) + outputs[1:]

       return TFBaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=outputs.last_hidden_state      ,  pooler_output=outputs.pooler_output      ,  hidden_states=outputs.hidden_states      ,  )






@add_start_docstrings(
    '''
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    ''', SCREAMING_SNAKE_CASE__, )
class      A     (   SCREAMING_SNAKE_CASE__, SCREAMING_SNAKE_CASE__    ):






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  *__lowerCAmelCase      ,  **__lowerCAmelCase			):
       super().__init__(__lowerCAmelCase      ,  *__lowerCAmelCase      ,  **__lowerCAmelCase			)
       UpperCamelCase_					: str            =				config.num_labels
       UpperCamelCase_					: Any            =				TFRegNetMainLayer(__lowerCAmelCase      ,  name="""regnet"""			)
       # classification head
       UpperCamelCase_					: List[Any]            =				[
           tf.keras.layers.Flatten(),
           tf.keras.layers.Dense(config.num_labels      ,  name="""classifier.1"""			) if config.num_labels > 0 else tf.identity,
       ]





 @unpack_inputs
 @add_start_docstrings_to_model_forward(__lowerCAmelCase			)
 @add_code_sample_docstrings(
     checkpoint=_IMAGE_CLASS_CHECKPOINT      ,  output_type=__lowerCAmelCase      ,  config_class=_CONFIG_FOR_DOC      ,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,  )
 def 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase = None      ,  __lowerCAmelCase=False      ,  ):
       UpperCamelCase_					: str            =				(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       UpperCamelCase_					: Tuple            =				return_dict if return_dict is not None else self.config.use_return_dict

       UpperCamelCase_					: Tuple            =				self.regnet(
           __lowerCAmelCase      ,  output_hidden_states=__lowerCAmelCase      ,  return_dict=__lowerCAmelCase      ,  training=__lowerCAmelCase			)

       UpperCamelCase_					: Any            =				outputs.pooler_output if return_dict else outputs[1]

       UpperCamelCase_					: str            =				self.classifier[0](__lowerCAmelCase			)
       UpperCamelCase_					: List[Any]            =				self.classifier[1](__lowerCAmelCase			)

       UpperCamelCase_					: int            =				None if labels is None else self.hf_compute_loss(labels=__lowerCAmelCase      ,  logits=__lowerCAmelCase			)

       if not return_dict:
             UpperCamelCase_					: Dict            =				(logits,) + outputs[2:]
             return ((loss,) + output) if loss is not None else output

       return TFSequenceClassifierOutput(loss=__lowerCAmelCase      ,  logits=__lowerCAmelCase      ,  hidden_states=outputs.hidden_states			)














'''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      A     :






 """simple docstring"""





 def __init__(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase=13      ,  __lowerCAmelCase=7      ,  __lowerCAmelCase=6      ,  __lowerCAmelCase=17      ,  __lowerCAmelCase=23      ,  __lowerCAmelCase=11      ,  __lowerCAmelCase=True      ,  ):
       UpperCamelCase_					: Tuple            =				parent
       UpperCamelCase_					: List[Any]            =				batch_size
       UpperCamelCase_					: List[str]            =				seq_length
       UpperCamelCase_					: Dict            =				act_dim
       UpperCamelCase_					: Optional[int]            =				state_dim
       UpperCamelCase_					: Dict            =				hidden_size
       UpperCamelCase_					: str            =				max_length
       UpperCamelCase_					: int            =				is_training
 def 	_UpperCAmelCase							(				self			):
       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_					: Union[str, Any]            =				floats_tensor((self.batch_size, self.seq_length, 1)			)
       UpperCamelCase_					: List[str]            =				floats_tensor((self.batch_size, self.seq_length, 1)			)
       UpperCamelCase_					: Optional[Any]            =				ids_tensor((self.batch_size, self.seq_length)      ,  vocab_size=10_00			)
       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 	_UpperCAmelCase							(				self			):
       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 	_UpperCAmelCase							(				self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  ):
       UpperCamelCase_					: Tuple            =				DecisionTransformerModel(config=__lowerCAmelCase			)
       model.to(__lowerCAmelCase			)
       model.eval()
       UpperCamelCase_					: str            =				model(__lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase			)

       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 	_UpperCAmelCase							(				self			):
       UpperCamelCase_					: Dict            =				self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase_
       )     ,      (
           UpperCamelCase_
       )     ,      (
           UpperCamelCase_
       )     ,      (
           UpperCamelCase_
       )     ,      (
           UpperCamelCase_
       )     ,      (
           UpperCamelCase_
       )     ,      (
           UpperCamelCase_
       )     ,      
       )					: str            =				config_and_inputs
       UpperCamelCase_					: 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      A     (   SCREAMING_SNAKE_CASE__, SCREAMING_SNAKE_CASE__, SCREAMING_SNAKE_CASE__, unittest.TestCase    ):






 """simple docstring"""





 __a      :    Tuple			  =    (DecisionTransformerModel,) if is_torch_available() else ()
 __a      :    Union[str, Any]			  =    ()
 __a      :    int			  =    {'''feature-extraction''': DecisionTransformerModel} if is_torch_available() else {}

 # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
 __a      :    Optional[Any]			  =    False

 # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
 __a      :    int			  =    False
 __a      :    Tuple			  =    False
 __a      :    Dict			  =    False
 __a      :    Union[str, Any]			  =    False
 __a      :    int			  =    False
 __a      :    List[Any]			  =    False
 __a      :    List[str]			  =    False
 __a      :    Dict			  =    False
 __a      :    Optional[Any]			  =    False
 def 	_UpperCAmelCase							(				self			):
       UpperCamelCase_					: Dict            =				DecisionTransformerModelTester(self			)
       UpperCamelCase_					: Dict            =				ConfigTester(self      ,  config_class=__lowerCAmelCase      ,  hidden_size=37			)
 def 	_UpperCAmelCase							(				self			):
       self.config_tester.run_common_tests()
 def 	_UpperCAmelCase							(				self			):
       UpperCamelCase_					: Optional[int]            =				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__lowerCAmelCase			)
 @slow
 def 	_UpperCAmelCase							(				self			):
       for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCamelCase_					: Union[str, Any]            =				DecisionTransformerModel.from_pretrained(__lowerCAmelCase			)
             self.assertIsNotNone(__lowerCAmelCase			)





 def 	_UpperCAmelCase							(				self			):
       UpperCamelCase_     ,      UpperCamelCase_					: str            =				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             UpperCamelCase_					: Union[str, Any]            =				model_class(__lowerCAmelCase			)
             UpperCamelCase_					: Optional[int]            =				inspect.signature(model.forward			)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCamelCase_					: Any            =				[*signature.parameters.keys()]

             UpperCamelCase_					: List[str]            =				[
                 """states""",
                 """actions""",
                 """rewards""",
                 """returns_to_go""",
                 """timesteps""",
                 """attention_mask""",
             ]

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






@require_torch
class      A     (   unittest.TestCase    ):






 """simple docstring"""





 @slow
 def 	_UpperCAmelCase							(				self			):
       UpperCamelCase_					: Tuple            =				2  # number of steps of autoregressive prediction we will perform
       UpperCamelCase_					: Dict            =				10  # defined by the RL environment, may be normalized
       UpperCamelCase_					: Any            =				DecisionTransformerModel.from_pretrained("""edbeeching/decision-transformer-gym-hopper-expert"""			)
       UpperCamelCase_					: Tuple            =				model.to(__lowerCAmelCase			)
       UpperCamelCase_					: List[str]            =				model.config
       torch.manual_seed(0			)
       UpperCamelCase_					: str            =				torch.randn(1      ,  1      ,  config.state_dim			).to(device=__lowerCAmelCase      ,  dtype=torch.floataa			)  # env.reset()

       UpperCamelCase_					: List[str]            =				torch.tensor(
           [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]]      ,  device=__lowerCAmelCase			)

       UpperCamelCase_					: Tuple            =				torch.tensor(__lowerCAmelCase      ,  device=__lowerCAmelCase      ,  dtype=torch.floataa			).reshape(1      ,  1      ,  1			)
       UpperCamelCase_					: Optional[int]            =				state
       UpperCamelCase_					: int            =				torch.zeros(1      ,  0      ,  config.act_dim      ,  device=__lowerCAmelCase      ,  dtype=torch.floataa			)
       UpperCamelCase_					: Optional[int]            =				torch.zeros(1      ,  0      ,  device=__lowerCAmelCase      ,  dtype=torch.floataa			)
       UpperCamelCase_					: Optional[Any]            =				torch.tensor(0      ,  device=__lowerCAmelCase      ,  dtype=torch.long			).reshape(1      ,  1			)

       for step in range(__lowerCAmelCase			):
             UpperCamelCase_					: List[Any]            =				torch.cat([actions, torch.zeros(1      ,  1      ,  config.act_dim      ,  device=__lowerCAmelCase			)]      ,  dim=1			)
             UpperCamelCase_					: Optional[int]            =				torch.cat([rewards, torch.zeros(1      ,  1      ,  device=__lowerCAmelCase			)]      ,  dim=1			)

             UpperCamelCase_					: Union[str, Any]            =				torch.ones(1      ,  states.shape[1]			).to(dtype=torch.long      ,  device=states.device			)

             with torch.no_grad():
                   UpperCamelCase_     ,      UpperCamelCase_     ,      UpperCamelCase_					: List[Any]            =				model(
                       states=__lowerCAmelCase      ,  actions=__lowerCAmelCase      ,  rewards=__lowerCAmelCase      ,  returns_to_go=__lowerCAmelCase      ,  timesteps=__lowerCAmelCase      ,  attention_mask=__lowerCAmelCase      ,  return_dict=__lowerCAmelCase      ,  )

             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_					: Union[str, Any]            =				(  # env.step(action)
                 torch.randn(1      ,  1      ,  config.state_dim			).to(device=__lowerCAmelCase      ,  dtype=torch.floataa			),
                 1.0,
                 False,
                 {},
             )

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






