


from ...configuration_utils import PretrainedConfig
from ...utils import logging


A  :     Union[str, Any]    					=    logging.get_logger(__name__)

A  :     Optional[int]    					=    {
    "edbeeching/decision-transformer-gym-hopper-medium": (
        "https://huggingface.co/edbeeching/decision-transformer-gym-hopper-medium/resolve/main/config.json"
    ),
    # See all DecisionTransformer models at https://huggingface.co/models?filter=decision_transformer
}
class 						_UpperCamelCase						(     lowerCAmelCase__	):
							'''simple docstring'''







							__UpperCAmelCase		:					Tuple          ="""decision_transformer"""
							__UpperCAmelCase		:					Any          =["""past_key_values"""]
							__UpperCAmelCase		:					List[str]          ={
							    """max_position_embeddings""": """n_positions""",
							    """num_attention_heads""": """n_head""",
							    """num_hidden_layers""": """n_layer""",
							}

							def __init__(  self					,     __a=17					,     __a=4					,     __a=1_28					,     __a=40_96					,     __a=True					,     __a=1					,     __a=10_24					,     __a=3					,     __a=1					,     __a=None					,     __a="relu"					,     __a=0.1					,     __a=0.1					,     __a=0.1					,     __a=1e-5					,     __a=0.0_2					,     __a=True					,     __a=True					,     __a=5_02_56					,     __a=5_02_56					,     __a=False					,     __a=False					,     **__a					,     ):
										__lowerCAmelCase       					=    state_dim
										__lowerCAmelCase       					=    act_dim
										__lowerCAmelCase       					=    hidden_size
										__lowerCAmelCase       					=    max_ep_len
										__lowerCAmelCase       					=    action_tanh
										__lowerCAmelCase       					=    vocab_size
										__lowerCAmelCase       					=    n_positions
										__lowerCAmelCase       					=    n_layer
										__lowerCAmelCase       					=    n_head
										__lowerCAmelCase       					=    n_inner
										__lowerCAmelCase       					=    activation_function
										__lowerCAmelCase       					=    resid_pdrop
										__lowerCAmelCase       					=    embd_pdrop
										__lowerCAmelCase       					=    attn_pdrop
										__lowerCAmelCase       					=    layer_norm_epsilon
										__lowerCAmelCase       					=    initializer_range
										__lowerCAmelCase       					=    scale_attn_weights
										__lowerCAmelCase       					=    use_cache
										__lowerCAmelCase       					=    scale_attn_by_inverse_layer_idx
										__lowerCAmelCase       					=    reorder_and_upcast_attn

										__lowerCAmelCase       					=    bos_token_id
										__lowerCAmelCase       					=    eos_token_id

										super().__init__(bos_token_id=__a					,     eos_token_id=__a					,     **__a			)








"""simple docstring"""


import flax.linen as nn
import jax
import jax.numpy as jnp
class 						_UpperCamelCase						(     nn.Module	):
							'''simple docstring'''







							__UpperCAmelCase		:					int
							__UpperCAmelCase		:					jnp.dtype          =jnp.floataa

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    nn.Conv(
										    self.out_channels					,     kernel_size=(3, 3)					,     strides=(1, 1)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )



							def __call__(  self					,     __a			):
										__lowerCAmelCase					,			__lowerCAmelCase					,			__lowerCAmelCase					,			__lowerCAmelCase       					=    hidden_states.shape
										__lowerCAmelCase       					=    jax.image.resize(
										    __a					,     shape=(batch, height * 2, width * 2, channels)					,     method="nearest"					,     )
										__lowerCAmelCase       					=    self.conv(__a			)
										return hidden_states
class 						_UpperCamelCase						(     nn.Module	):
							'''simple docstring'''







							__UpperCAmelCase		:					int
							__UpperCAmelCase		:					jnp.dtype          =jnp.floataa

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    nn.Conv(
										    self.out_channels					,     kernel_size=(3, 3)					,     strides=(2, 2)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )



							def __call__(  self					,     __a			):
										# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
										# hidden_states = jnp.pad(hidden_states, pad_width=pad)
										__lowerCAmelCase       					=    self.conv(__a			)
										return hidden_states
class 						_UpperCamelCase						(     nn.Module	):
							'''simple docstring'''







							__UpperCAmelCase		:					int
							__UpperCAmelCase		:					int          =None
							__UpperCAmelCase		:					float          =0.0
							__UpperCAmelCase		:					bool          =None
							__UpperCAmelCase		:					jnp.dtype          =jnp.floataa

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.in_channels if self.out_channels is None else self.out_channels

										__lowerCAmelCase       					=    nn.GroupNorm(num_groups=32					,     epsilon=1e-5			)
										__lowerCAmelCase       					=    nn.Conv(
										    __a					,     kernel_size=(3, 3)					,     strides=(1, 1)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )

										__lowerCAmelCase       					=    nn.Dense(__a					,     dtype=self.dtype			)

										__lowerCAmelCase       					=    nn.GroupNorm(num_groups=32					,     epsilon=1e-5			)
										__lowerCAmelCase       					=    nn.Dropout(self.dropout_prob			)
										__lowerCAmelCase       					=    nn.Conv(
										    __a					,     kernel_size=(3, 3)					,     strides=(1, 1)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )

										__lowerCAmelCase       					=    self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

										__lowerCAmelCase       					=    None
										if use_nin_shortcut:
													__lowerCAmelCase       					=    nn.Conv(
													    __a					,     kernel_size=(1, 1)					,     strides=(1, 1)					,     padding="VALID"					,     dtype=self.dtype					,     )



							def __call__(  self					,     __a					,     __a					,     __a=True			):
										__lowerCAmelCase       					=    hidden_states
										__lowerCAmelCase       					=    self.norma(__a			)
										__lowerCAmelCase       					=    nn.swish(__a			)
										__lowerCAmelCase       					=    self.conva(__a			)

										__lowerCAmelCase       					=    self.time_emb_proj(nn.swish(__a			)			)
										__lowerCAmelCase       					=    jnp.expand_dims(jnp.expand_dims(__a					,     1			)					,     1			)
										__lowerCAmelCase       					=    hidden_states + temb

										__lowerCAmelCase       					=    self.norma(__a			)
										__lowerCAmelCase       					=    nn.swish(__a			)
										__lowerCAmelCase       					=    self.dropout(__a					,     __a			)
										__lowerCAmelCase       					=    self.conva(__a			)

										if self.conv_shortcut is not None:
													__lowerCAmelCase       					=    self.conv_shortcut(__a			)

										return hidden_states + residual




