docstr: Args:
      x: shape (batch, N, 9)
      adj: shape (batch, 4, N, N)
      x_cont: shape (batch, 4)
      index: shape (batch, 2)
  Returns:
      x_cont: shape (batch, 4)
      x_log_jacob: shape (batch, )            
docstr_tok: ["Args",":","x",":","shape","(","batch",",","N",",","9",")","adj",":","shape","(","batch",",","4",",","N",",","N",")","x_cont",":","shape","(","batch",",","4",")","index",":","shape","(","batch",",","2",")","Returns",":","x_cont",":","shape","(","batch",",","4",")","x_log_jacob",":","shape","(","batch",",",")"]
code: def forward_rl_edge(self, x, adj, x_cont, index):
      """
      Args:
          x: shape (batch, N, 9)
          adj: shape (batch, 4, N, N)
          x_cont: shape (batch, 4)
          index: shape (batch, 2)
      Returns:
          x_cont: shape (batch, 4)
          x_log_jacob: shape (batch, )
      """
      embs = self._get_embs_edge(x, adj, index)  # (batch, 3d)
  
      for i in range(self.num_flow_layer):
          edge_s, edge_t = self.edge_st_net[i](embs)
  
          if self.st_type == "sigmoid":
              x_cont = x_cont * edge_s + edge_t
          elif self.st_type == "exp":
              edge_s = edge_s.exp()
              x_cont = (x_cont + edge_t) * edge_s
          elif self.st_type == "softplus":
              x_cont = (x_cont + edge_t) * edge_s
          else:
              raise ValueError(
                  "unsupported st type: (%s)" % self.args.st_type
              )
  
          if torch.isnan(x_cont).any():
              raise RuntimeError(
                  "x_cont has NaN entries after transformation at layer %d"
                  % i
              )
  
          if i == 0:
              x_log_jacob = (torch.abs(edge_s) + 1e-20).log()
          else:
              x_log_jacob += (torch.abs(edge_s) + 1e-20).log()
  
      x_log_jacob = x_log_jacob.sum(-1)  # (batch)
  
      return x_cont, x_log_jacob
  
file_path: DIG/dig/ggraph/method/GraphAF/model/graphaf.py
imports: import numpy
  import torch
  import torch.nn
  from .rgcn import *
  from .st_net import *
label_desc: Use image node_id 4 for calling the MaskedGraphAF obj's underlying member method code with example usage: obj.forward_rl_edge(x, adj, x_cont, index) and returns: x_cont,  x_log_jacob