class_docstr: Transformer with extra features.
  
  This module is the backbone of `FTTransformer`.
class_docstr_tok: ["Transformer","with","extra","features",".","This","module","is","the","backbone","of","`","FTTransformer","`","."]
code: def forward(self, x: Tensor) -> Tensor:
      assert (
          x.ndim == 3
      ), "The input must have 3 dimensions: (n_objects, n_tokens, d_token)"
      for layer_idx, layer in enumerate(self.blocks):
          layer = cast(nn.ModuleDict, layer)
  
          query_idx = (
              self.last_layer_query_idx
              if layer_idx + 1 == len(self.blocks)
              else None
          )
          x_residual = self._start_residual(layer, "attention", x)
          x_residual, _ = layer["attention"](
              x_residual
              if query_idx is None
              else x_residual[:, query_idx],
              x_residual,
              *self._get_kv_compressions(layer),
          )
          if query_idx is not None:
              x = x[:, query_idx]
          x = self._end_residual(layer, "attention", x, x_residual)
  
          x_residual = self._start_residual(layer, "ffn", x)
          x_residual = layer["ffn"](x_residual)
          x = self._end_residual(layer, "ffn", x, x_residual)
          x = layer["output"](x)
  
      x = self.head(x)
  
      return x
  
imports: import enum
  import math
  import warnings
  from typing import Callable, Dict, List, Optional, Tuple, Union, cast
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from .utils import init_weights
label_desc: Use image node_id 5 for calling the Custom_Transformer obj's underlying member method code with example usage: obj.forward(x) and returns: x