class_docstr: An implementation of the tranform attention mechanism.
  For details see this paper: `"GMAN: A Graph Multi-Attention Network for Traffic Prediction."
  <https://arxiv.org/pdf/1911.08415.pdf>`_
  
  Args:
      K (int) : Number of attention heads.
      d (int) : Dimension of each attention head outputs.
      bn_decay (float): Batch normalization momentum.
docstr: Making a forward pass of the transform attention layer.
  
  Arg types:
      * **X** (PyTorch Float Tensor) - Input sequence, with shape (batch_size, num_his, num_nodes, K*d).
      * **STE_his** (Pytorch Float Tensor) - Spatial-temporal embedding for history,
      with shape (batch_size, num_his, num_nodes, K*d).
      * **STE_pred** (Pytorch Float Tensor) - Spatial-temporal embedding for prediction,
      with shape (batch_size, num_pred, num_nodes, K*d).
  
  Return types:
      * **X** (PyTorch Float Tensor) - Output sequence for prediction, with shape (batch_size, num_pred, num_nodes, K*d).
docstr_tok: ["Making","a","forward","pass","of","the","transform","attention","layer",".","Arg","types",":","*","*","*","X","*","*","(","PyTorch","Float","Tensor",")","-","Input","sequence",",","with","shape","(","batch_size",",","num_his",",","num_nodes",",","K","*","d",")",".","*","*","*","STE_his","*","*","(","Pytorch","Float","Tensor",")","-","Spatial-temporal","embedding","for","history",",","with","shape","(","batch_size",",","num_his",",","num_nodes",",","K","*","d",")",".","*","*","*","STE_pred","*","*","(","Pytorch","Float","Tensor",")","-","Spatial-temporal","embedding","for","prediction",",","with","shape","(","batch_size",",","num_pred",",","num_nodes",",","K","*","d",")",".","Return","types",":","*","*","*","X","*","*","(","PyTorch","Float","Tensor",")","-","Output","sequence","for","prediction",",","with","shape","(","batch_size",",","num_pred",",","num_nodes",",","K","*","d",")","."]
code: def forward(
      self,
      X: torch.FloatTensor,
      STE_his: torch.FloatTensor,
      STE_pred: torch.FloatTensor,
  ) -> torch.FloatTensor:
      """
      Making a forward pass of the transform attention layer.
  
      Arg types:
          * **X** (PyTorch Float Tensor) - Input sequence, with shape (batch_size, num_his, num_nodes, K*d).
          * **STE_his** (Pytorch Float Tensor) - Spatial-temporal embedding for history,
          with shape (batch_size, num_his, num_nodes, K*d).
          * **STE_pred** (Pytorch Float Tensor) - Spatial-temporal embedding for prediction,
          with shape (batch_size, num_pred, num_nodes, K*d).
  
      Return types:
          * **X** (PyTorch Float Tensor) - Output sequence for prediction, with shape (batch_size, num_pred, num_nodes, K*d).
      """
      batch_size = X.shape[0]
      query = self._fully_connected_q(STE_pred)
      key = self._fully_connected_k(STE_his)
      value = self._fully_connected_v(X)
      query = torch.cat(torch.split(query, self._K, dim=-1), dim=0)
      key = torch.cat(torch.split(key, self._K, dim=-1), dim=0)
      value = torch.cat(torch.split(value, self._K, dim=-1), dim=0)
      query = query.permute(0, 2, 1, 3)
      key = key.permute(0, 2, 3, 1)
      value = value.permute(0, 2, 1, 3)
      attention = torch.matmul(query, key)
      attention /= self._d**0.5
      attention = F.softmax(attention, dim=-1)
      X = torch.matmul(attention, value)
      X = X.permute(0, 2, 1, 3)
      X = torch.cat(torch.split(X, batch_size, dim=0), dim=-1)
      X = self._fully_connected(X)
      del query, key, value, attention
      return X
  
imports: import math
  from typing import Union, Callable, Optional
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the TransformAttention obj's underlying member method code with example usage: obj.forward(X, STE_his, STE_pred) and returns: X