class_docstr: Torso model of OpenAlphaTensor.
  
  It maps an input tensor of shape (N, T, S, S, S) to (N, 3S*S, c), where:
  
      N is the batch size;
      T is the context size (size of the history + 1);
      S is the number of elements in each matrix to be multiplied;
      c is the output dimensionality.
code: def forward(self, x: torch.Tensor, scalars: torch.Tensor):
      # x.size = (N, T, S, S, S)
      # scalars.size = (N, s)
      batch_size = x.shape[0]
      S = x.shape[-1]
      T = x.shape[1]
      x1 = x.permute(0, 2, 3, 4, 1).reshape(batch_size, S, S, S * T)
      x2 = x.permute(0, 4, 2, 3, 1).reshape(batch_size, S, S, S * T)
      x3 = x.permute(0, 3, 4, 2, 1).reshape(batch_size, S, S, S * T)
      input_list = [x1, x2, x3]
      for i in range(3):
          temp = self.linears_1[i](scalars).reshape(batch_size, S, S, 1)
          input_list[i] = torch.cat([input_list[i], temp], dim=-1)
          input_list[i] = self.linears_2[i](input_list[i])
      x1, x2, x3 = input_list
      for layer in self.attentive_modes:
          x1, x2, x3 = layer(x1, x2, x3)
      return torch.stack([x1, x2, x3], dim=2).reshape(
          batch_size, 3 * S * S, -1
      )
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/modules/torso.py
imports: import torch
  from open_alpha_tensor.core.modules.attention import AlphaMultiHeadAttention
label_desc: Use image node_id 2 for calling the TorsoModel obj's underlying member method code with example usage: obj.forward(x, scalars) and returns: torch