class_docstr: Equivariant Graph Convolutional Layer from `E(n) Equivariant Graph
  Neural Networks <https://arxiv.org/abs/2102.09844>`__
  
  .. math::
  
      m_{ij}=\phi_e(h_i^l, h_j^l, ||x_i^l-x_j^l||^2, a_{ij})
  
      x_i^{l+1} = x_i^l + C\sum_{j\in\mathcal{N}(i)}(x_i^l-x_j^l)\phi_x(m_{ij})
  
      m_i = \sum_{j\in\mathcal{N}(i)} m_{ij}
  
      h_i^{l+1} = \phi_h(h_i^l, m_i)
  
  where :math:`h_i`, :math:`x_i`, :math:`a_{ij}` are node features, coordinate
  features, and edge features respectively. :math:`\phi_e`, :math:`\phi_h`, and
  :math:`\phi_x` are two-layer MLPs. :math:`C` is a constant for normalization,
  computed as :math:`1/|\mathcal{N}(i)|`.
  
  Parameters
  ----------
  in_size : int
      Input feature size; i.e. the size of :math:`h_i^l`.
  hidden_size : int
      Hidden feature size; i.e. the size of hidden layer in the two-layer MLPs in
      :math:`\phi_e, \phi_x, \phi_h`.
  out_size : int
      Output feature size; i.e. the size of :math:`h_i^{l+1}`.
  edge_feat_size : int, optional
      Edge feature size; i.e. the size of :math:`a_{ij}`. Default: 0.
  
  Example
  -------
  >>> import dgl
  >>> import torch as th
  >>> from dgl.nn import EGNNConv
  >>>
  >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
  >>> node_feat, coord_feat, edge_feat = th.ones(6, 10), th.ones(6, 3), th.ones(6, 2)
  >>> conv = EGNNConv(10, 10, 10, 2)
  >>> h, x = conv(g, node_feat, coord_feat, edge_feat)
docstr: Description
  -----------
  Compute EGNN layer.
  
  Parameters
  ----------
  graph : DGLGraph
      The graph.
  node_feat : torch.Tensor
      The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
      nodes, and :math:`h_n` must be the same as in_size.
  coord_feat : torch.Tensor
      The coordinate feature of shape :math:`(N, h_x)`. :math:`N` is the
      number of nodes, and :math:`h_x` can be any positive integer.
  edge_feat : torch.Tensor, optional
      The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
      edges, and :math:`h_e` must be the same as edge_feat_size.
  
  Returns
  -------
  node_feat_out : torch.Tensor
      The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
      is the same as out_size.
  coord_feat_out: torch.Tensor
      The output coordinate feature of shape :math:`(N, h_x)` where :math:`h_x`
      is the same as the input coordinate feature dimension.
docstr_tok: ["Description","--","--","--","--","--","-","Compute","EGNN","layer",".","Parameters","--","--","--","--","--","graph",":","DGLGraph","The","graph",".","node_feat",":","torch.Tensor","The","input","feature","of","shape",":","math",":","`","(","N",",","h_n",")","`",".",":","math",":","`","N","`","is","the","number","of","nodes",",","and",":","math",":","`","h_n","`","must","be","the","same","as","in_size",".","coord_feat",":","torch.Tensor","The","coordinate","feature","of","shape",":","math",":","`","(","N",",","h_x",")","`",".",":","math",":","`","N","`","is","the","number","of","nodes",",","and",":","math",":","`","h_x","`","can","be","any","positive","integer",".","edge_feat",":","torch.Tensor",",","optional","The","edge","feature","of","shape",":","math",":","`","(","M",",","h_e",")","`",".",":","math",":","`","M","`","is","the","number","of","edges",",","and",":","math",":","`","h_e","`","must","be","the","same","as","edge_feat_size",".","Returns","--","--","--","-","node_feat_out",":","torch.Tensor","The","output","node","feature","of","shape",":","math",":","`","(","N",",","h_n","'",")","`","where",":","math",":","`","h_n","'","`","is","the","same","as","out_size",".","coord_feat_out",":","torch.Tensor","The","output","coordinate","feature","of","shape",":","math",":","`","(","N",",","h_x",")","`","where",":","math",":","`","h_x","`","is","the","same","as","the","input","coordinate","feature","dimension","."]
code: def forward(self, graph, node_feat, coord_feat, edge_feat=None):
      r"""
      Description
      -----------
      Compute EGNN layer.
  
      Parameters
      ----------
      graph : DGLGraph
          The graph.
      node_feat : torch.Tensor
          The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
          nodes, and :math:`h_n` must be the same as in_size.
      coord_feat : torch.Tensor
          The coordinate feature of shape :math:`(N, h_x)`. :math:`N` is the
          number of nodes, and :math:`h_x` can be any positive integer.
      edge_feat : torch.Tensor, optional
          The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
          edges, and :math:`h_e` must be the same as edge_feat_size.
  
      Returns
      -------
      node_feat_out : torch.Tensor
          The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
          is the same as out_size.
      coord_feat_out: torch.Tensor
          The output coordinate feature of shape :math:`(N, h_x)` where :math:`h_x`
          is the same as the input coordinate feature dimension.
      """
      with graph.local_scope():
          # node feature
          graph.ndata["h"] = node_feat
          # coordinate feature
          graph.ndata["x"] = coord_feat
          # edge feature
          if self.edge_feat_size > 0:
              assert (
                  edge_feat is not None
              ), "Edge features must be provided."
              graph.edata["a"] = edge_feat
          # get coordinate diff & radial features
          graph.apply_edges(fn.u_sub_v("x", "x", "x_diff"))
          graph.edata["radial"] = (
              graph.edata["x_diff"].square().sum(dim=1).unsqueeze(-1)
          )
          # normalize coordinate difference
          graph.edata["x_diff"] = graph.edata["x_diff"] / (
              graph.edata["radial"].sqrt() + 1e-30
          )
          graph.apply_edges(self.message)
          graph.update_all(
              fn.copy_e("msg_x", "m"), fn.mean("m", "x_neigh")
          )
          graph.update_all(
              fn.copy_e("msg_h", "m"), fn.sum("m", "h_neigh")
          )
  
          h_neigh, x_neigh = (
              graph.ndata["h_neigh"],
              graph.ndata["x_neigh"],
          )
  
          h = self.node_mlp(torch.cat([node_feat, h_neigh], dim=-1))
          x = coord_feat + x_neigh
  
          return h, x
  
file_path: dgl/python/dgl/nn/pytorch/conv/egnnconv.py
imports: import torch
  import torch.nn
  from ....None import function
label_desc: Use image node_id 3 for calling the EGNNConv obj's underlying member method code with example usage: obj.forward(graph, node_feat, coord_feat, edge_feat) and returns: h,  x