docstr: :param x of shape (batch, N, 9)
  :param adj of shape (batch, 4, N, N)
  :return: inputs for st_net_node and st_net_edge
  graph_emb_node of shape (batch, N, d)
  graph_emb_edge of shape (batch, repeat-N, 3d)
docstr_tok: [":","param","x","of","shape","(","batch",",","N",",","9",")",":","param","adj","of","shape","(","batch",",","4",",","N",",","N",")",":","return",":","inputs","for","st_net_node","and","st_net_edge","graph_emb_node","of","shape","(","batch",",","N",",","d",")","graph_emb_edge","of","shape","(","batch",",","repeat-N",",","3d",")"]
code: def _get_embs(self, x, adj):
      """
      :param x of shape (batch, N, 9)
      :param adj of shape (batch, 4, N, N)
      :return: inputs for st_net_node and st_net_edge
      graph_emb_node of shape (batch, N, d)
      graph_emb_edge of shape (batch, repeat-N, 3d)
  
      """
      # inputs for RelGCNs
      batch_size = x.size(0)
      adj = adj[
          :, :3
      ]  # (batch, 3, N, N) TODO: check whether we have to use the 4-th slices(virtual bond) or not
      x = torch.where(
          self.mask_node,
          x.unsqueeze(1).repeat(1, self.repeat_num, 1, 1),
          torch.zeros([1], device=x.device),
      ).view(
          -1, self.graph_size, self.num_node_type
      )  # (batch*repeat_num, N, 9)
  
      adj = torch.where(
          self.mask_edge,
          adj.unsqueeze(1).repeat(1, self.repeat_num, 1, 1, 1),
          torch.zeros([1], device=x.device),
      ).view(
          -1, self.num_edge_type - 1, self.graph_size, self.graph_size
      )  # (batch*repeat_num, 3, N, N)
      node_emb = self.rgcn(x, adj)  # (batch*repeat_num, N, d)
  
      if hasattr(self, "batchNorm"):
          node_emb = self.batchNorm(node_emb.transpose(1, 2)).transpose(
              1, 2
          )  # (batch*repeat_num, N, d)
  
      node_emb = node_emb.view(
          batch_size, self.repeat_num, self.graph_size, -1
      )  # (batch, repeat_num, N, d)
  
      graph_emb = torch.sum(
          node_emb, dim=2, keepdim=False
      )  # (batch, repeat_num, d)
  
      #  input for st_net_node
      graph_emb_node = graph_emb[
          :, : self.graph_size
      ].contiguous()  # (batch, N, d)
      # graph_emb_node = graph_emb_node.view(batch_size * self.graph_size, -1)  # (batch*N, d)
  
      # input for st_net_edge
      graph_emb_edge = graph_emb[
          :, self.graph_size :
      ].contiguous()  # (batch, repeat_num-N, d)
      graph_emb_edge = graph_emb_edge.unsqueeze(
          2
      )  # (batch, repeat_num-N, 1, d)
  
      all_node_emb_edge = node_emb[
          :, self.graph_size :
      ]  # (batch, repeat_num-N, N, d)
  
      index = self.index_select_edge.view(1, -1, 2, 1).repeat(
          batch_size, 1, 1, self.emb_size
      )  # (batch_size, repeat_num-N, 2, d)
  
      graph_node_emb_edge = torch.cat(
          (
              torch.gather(all_node_emb_edge, dim=2, index=index),
              graph_emb_edge,
          ),
          dim=2,
      )  # (batch_size, repeat_num-N, 3, d)
  
      graph_node_emb_edge = graph_node_emb_edge.view(
          batch_size, self.repeat_num - self.graph_size, -1
      )  # (batch_size, (repeat_num-N), 3*d)
  
      return graph_emb_node, graph_node_emb_edge
  
file_path: DIG/dig/ggraph/method/GraphDF/model/disgraphaf.py
imports: import torch
  from torch import nn
  from .st_net import ST_Dis
  from .df_utils import one_hot_add, one_hot_minus
  from .rgcn import RGCN
label_desc: Use image node_id 8 for calling the DisGraphAF obj's underlying member method code with example usage: obj._get_embs(x, adj) and returns: graph_emb_node,  graph_node_emb_edge