class_docstr: An implementation of the Attention Based Spatial-Temporal Graph Convolutional Block.
  For details see this paper: `"Attention Based Spatial-Temporal Graph Convolutional
  Networks for Traffic Flow Forecasting." <https://ojs.aaai.org/index.php/AAAI/article/view/3881>`_
  
  Args:
      in_channels (int): Number of input features.
      K (int): Order of Chebyshev polynomials. Degree is K-1.
      nb_chev_filter (int): Number of Chebyshev filters.
      nb_time_filter (int): Number of time filters.
      time_strides (int): Time strides during temporal convolution.
      num_of_vertices (int): Number of vertices in the graph.
      num_of_timesteps (int): Number of time lags.
      normalization (str, optional): The normalization scheme for the graph
          Laplacian (default: :obj:`"sym"`):
          1. :obj:`None`: No normalization
          :math:`\mathbf{L} = \mathbf{D} - \mathbf{A}`
          2. :obj:`"sym"`: Symmetric normalization
          :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1/2} \mathbf{A}
          \mathbf{D}^{-1/2}`
          3. :obj:`"rw"`: Random-walk normalization
          :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1} \mathbf{A}`
          You need to pass :obj:`lambda_max` to the :meth:`forward` method of
          this operator in case the normalization is non-symmetric.
          :obj:`\lambda_max` should be a :class:`torch.Tensor` of size
          :obj:`[num_graphs]` in a mini-batch scenario and a
          scalar/zero-dimensional tensor when operating on single graphs.
          You can pre-compute :obj:`lambda_max` via the
          :class:`torch_geometric.transforms.LaplacianLambdaMax` transform.
      bias (bool, optional): If set to :obj:`False`, the layer will not learn
          an additive bias. (default: :obj:`True`)
docstr: Making a forward pass with the ASTGCN block.
  
  Arg types:
      * **X** (PyTorch Float Tensor) - Node features for T time periods, with shape (B, N_nodes, F_in, T_in).
      * **edge_index** (LongTensor): Edge indices, can be an array of a list of Tensor arrays, depending on whether edges change over time.
  
  Return types:
      * **X** (PyTorch Float Tensor) - Hidden state tensor for all nodes, with shape (B, N_nodes, nb_time_filter, T_out).
code: def forward(
      self,
      X: torch.FloatTensor,
      edge_index: Union[torch.LongTensor, List[torch.LongTensor]],
  ) -> torch.FloatTensor:
      """
      Making a forward pass with the ASTGCN block.
  
      Arg types:
          * **X** (PyTorch Float Tensor) - Node features for T time periods, with shape (B, N_nodes, F_in, T_in).
          * **edge_index** (LongTensor): Edge indices, can be an array of a list of Tensor arrays, depending on whether edges change over time.
  
      Return types:
          * **X** (PyTorch Float Tensor) - Hidden state tensor for all nodes, with shape (B, N_nodes, nb_time_filter, T_out).
      """
      (
          batch_size,
          num_of_vertices,
          num_of_features,
          num_of_timesteps,
      ) = X.shape  # (32, 307, 1, 12)
  
      X_tilde = self._temporal_attention(
          X
      )  # (b, T, T)  (32, 12, 12) * reshaped x(32, 307, 12)  -reshape> (32, 307, 1, 12)
      # xreshaped is e.g. (32, 307, 12) * (32, 12, 12) -then_reshaped> (32, 307, 1, 12)
      X_tilde = torch.matmul(
          X.reshape(batch_size, -1, num_of_timesteps), X_tilde
      )
      X_tilde = X_tilde.reshape(
          batch_size, num_of_vertices, num_of_features, num_of_timesteps
      )
      X_tilde = self._spatial_attention(
          X_tilde
      )  # (B,N,N) for example (32, 307, 307)
  
      if not isinstance(edge_index, list):
          data = Data(
              edge_index=edge_index,
              edge_attr=None,
              num_nodes=num_of_vertices,
          )
          if self._normalization != "sym":
              lambda_max = LaplacianLambdaMax()(data).lambda_max
          else:
              lambda_max = None
          X_hat = []
          for t in range(num_of_timesteps):
              X_hat.append(
                  torch.unsqueeze(
                      self._chebconv_attention(
                          X[:, :, :, t],
                          edge_index,
                          X_tilde,
                          lambda_max=lambda_max,
                      ),
                      -1,
                  )
              )
  
          X_hat = F.relu(torch.cat(X_hat, dim=-1))
      else:
          X_hat = []
          for t in range(num_of_timesteps):
              data = Data(
                  edge_index=edge_index[t],
                  edge_attr=None,
                  num_nodes=num_of_vertices,
              )
              if self._normalization != "sym":
                  lambda_max = LaplacianLambdaMax()(data).lambda_max
              else:
                  lambda_max = None
              X_hat.append(
                  torch.unsqueeze(
                      self._chebconv_attention(
                          X[:, :, :, t],
                          edge_index[t],
                          X_tilde,
                          lambda_max=lambda_max,
                      ),
                      -1,
                  )
              )
          X_hat = F.relu(torch.cat(X_hat, dim=-1))
  
      # (b,N,F,T)->(b,F,N,T) for example (32, 307, 64, 12) -premute->(32, 64, 307,12)
      # then convolution along the time axis is applied
      X_hat = self._time_convolution(
          X_hat.permute(0, 2, 1, 3)
      )  # will give (32, 64, 307,12)
      # (b,N,F,T)-permute>(b,F,N,T) (1,1)->(b,F,N,T)  (32, 64, 307, 12)
      X = self._residual_convolution(
          X.permute(0, 2, 1, 3)
      )  # will also give (32, 64, 307,12)
      # -adding X + X_hat->(32, 64, 307, 12)-premuting-> (32, 12, 307, 64)-layer_normalization_-premuting->(32, 307, 64,12)
      X = self._layer_norm(F.relu(X + X_hat).permute(0, 3, 2, 1))
      X = X.permute(0, 2, 3, 1)
      return X
  
imports: import math
  from typing import Optional, List, Union
  import torch
  import torch.nn
  from torch.nn import Parameter
  import torch.nn.functional
  from torch_geometric.data import Data
  from torch_geometric.typing import OptTensor
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.transforms import LaplacianLambdaMax
  from torch_geometric.utils import remove_self_loops, add_self_loops, get_laplacian
label_desc: Use image node_id 3 for calling the ASTGCNBlock obj's underlying member method code with example usage: obj.forward(X, edge_index) and returns: X