class_docstr: GTrXL wrapper serving as interface for ModelV2s that set use_attention.
class_docstr_tok: ["GTrXL","wrapper","serving","as","interface","for","ModelV2s","that","set","use_attention","."]
code: def forward(
      self,
      input_dict: Dict[str, TensorType],
      state: List[TensorType],
      seq_lens: TensorType,
  ) -> (TensorType, List[TensorType]):
      assert seq_lens is not None
      # Push obs through "unwrapped" net's `forward()` first.
      wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
  
      # Concat. prev-action/reward if required.
      prev_a_r = []
  
      # Prev actions.
      if self.use_n_prev_actions:
          prev_n_actions = input_dict[SampleBatch.PREV_ACTIONS]
          # If actions are not processed yet (in their original form as
          # have been sent to environment):
          # Flatten/one-hot into 1D array.
          if self.model_config["_disable_action_flattening"]:
              # Merge prev n actions into flat tensor.
              flat = flatten_inputs_to_1d_tensor(
                  prev_n_actions,
                  spaces_struct=self.action_space_struct,
                  time_axis=True,
              )
              # Fold time-axis into flattened data.
              flat = torch.reshape(flat, [flat.shape[0], -1])
              prev_a_r.append(flat)
          # If actions are already flattened (but not one-hot'd yet!),
          # one-hot discrete/multi-discrete actions here and concatenate the
          # n most recent actions together.
          else:
              if isinstance(self.action_space, Discrete):
                  for i in range(self.use_n_prev_actions):
                      prev_a_r.append(
                          one_hot(
                              prev_n_actions[:, i].float(),
                              space=self.action_space,
                          )
                      )
              elif isinstance(self.action_space, MultiDiscrete):
                  for i in range(
                      0,
                      self.use_n_prev_actions,
                      self.action_space.shape[0],
                  ):
                      prev_a_r.append(
                          one_hot(
                              prev_n_actions[
                                  :, i : i + self.action_space.shape[0]
                              ].float(),
                              space=self.action_space,
                          )
                      )
              else:
                  prev_a_r.append(
                      torch.reshape(
                          prev_n_actions.float(),
                          [
                              -1,
                              self.use_n_prev_actions * self.action_dim,
                          ],
                      )
                  )
      # Prev rewards.
      if self.use_n_prev_rewards:
          prev_a_r.append(
              torch.reshape(
                  input_dict[SampleBatch.PREV_REWARDS].float(),
                  [-1, self.use_n_prev_rewards],
              )
          )
  
      # Concat prev. actions + rewards to the "main" input.
      if prev_a_r:
          wrapped_out = torch.cat([wrapped_out] + prev_a_r, dim=1)
  
      # Then through our GTrXL.
      input_dict["obs_flat"] = input_dict["obs"] = wrapped_out
  
      self._features, memory_outs = self.gtrxl(
          input_dict, state, seq_lens
      )
      model_out = self._logits_branch(self._features)
      return model_out, memory_outs
  
imports: import gymnasium
  from gymnasium.spaces import Box, Discrete, MultiDiscrete
  import numpy
  import tree
  from typing import Dict, Optional, Union
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.misc import SlimFC
  from ray.rllib.models.torch.modules import GRUGate, RelativeMultiHeadAttention, SkipConnection
  from ray.rllib.models.torch.recurrent_net import RecurrentNetwork
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space
  from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor, one_hot
  from ray.rllib.utils.typing import ModelConfigDict, TensorType, List
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.util import log_once
label_desc: Use image node_id 2 for calling the AttentionWrapper obj's underlying member method code with example usage: obj.forward(input_dict, state, seq_lens) and returns: model_out,  memory_outs