class_docstr: The re-implementation for DimeNet++ from the `"Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules" <https://arxiv.org/abs/2011.14115>`_ paper
  under the 3DGN gramework from `"Spherical Message Passing for 3D Molecular Graphs" <https://openreview.net/forum?id=givsRXsOt9r>`_ paper.
  
  Args:
      energy_and_force (bool, optional): If set to :obj:`True`, will predict energy and take the negative of the derivative of the energy with respect to the atomic positions as predicted forces. (default: :obj:`False`)
      cutoff (float, optional): Cutoff distance for interatomic interactions. (default: :obj:`5.0`)
      num_layers (int, optional): Number of building blocks. (default: :obj:`4`)
      hidden_channels (int, optional): Hidden embedding size. (default: :obj:`128`)
      out_channels (int, optional): Size of each output sample. (default: :obj:`1`)
      int_emb_size (int, optional): Embedding size used for interaction triplets. (default: :obj:`64`)
      basis_emb_size (int, optional): Embedding size used in the basis transformation. (default: :obj:`8`)
      out_emb_channels (int, optional): Embedding size used for atoms in the output block. (default: :obj:`256`)
      num_spherical (int, optional): Number of spherical harmonics. (default: :obj:`7`)
      num_radial (int, optional): Number of radial basis functions. (default: :obj:`6`)
      envelop_exponent (int, optional): Shape of the smooth cutoff. (default: :obj:`5`)
      num_before_skip (int, optional): Number of residual layers in the interaction blocks before the skip connection. (default: :obj:`1`)
      num_after_skip (int, optional): Number of residual layers in the interaction blocks before the skip connection. (default: :obj:`2`)
      num_output_layers (int, optional): Number of linear layers for the output blocks. (default: :obj:`3`)
      act: (function, optional): The activation funtion. (default: :obj:`swish`) 
      output_init: (str, optional): The initialization fot the output. It could be :obj:`GlorotOrthogonal` and :obj:`zeros`. (default: :obj:`GlorotOrthogonal`)       
class_docstr_tok: ["The","re-implementation","for","DimeNet++","from","the","`","``","Fast","and","Uncertainty-Aware","Directional","Message","Passing","for","Non-Equilibrium","Molecules","''","<","https",":","\/\/arxiv.org\/abs\/2011.14115",">","`","_","paper","under","the","3DGN","gramework","from","`","``","Spherical","Message","Passing","for","3D","Molecular","Graphs","''","<","https",":","\/\/openreview.net\/forum","?","id=givsRXsOt9r",">","`","_","paper",".","Args",":","energy_and_force","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","will","predict","energy","and","take","the","negative","of","the","derivative","of","the","energy","with","respect","to","the","atomic","positions","as","predicted","forces",".","(","default",":",":","obj",":","`","False","`",")","cutoff","(","float",",","optional",")",":","Cutoff","distance","for","interatomic","interactions",".","(","default",":",":","obj",":","`","5.0","`",")","num_layers","(","int",",","optional",")",":","Number","of","building","blocks",".","(","default",":",":","obj",":","`","4","`",")","hidden_channels","(","int",",","optional",")",":","Hidden","embedding","size",".","(","default",":",":","obj",":","`","128","`",")","out_channels","(","int",",","optional",")",":","Size","of","each","output","sample",".","(","default",":",":","obj",":","`","1","`",")","int_emb_size","(","int",",","optional",")",":","Embedding","size","used","for","interaction","triplets",".","(","default",":",":","obj",":","`","64","`",")","basis_emb_size","(","int",",","optional",")",":","Embedding","size","used","in","the","basis","transformation",".","(","default",":",":","obj",":","`","8","`",")","out_emb_channels","(","int",",","optional",")",":","Embedding","size","used","for","atoms","in","the","output","block",".","(","default",":",":","obj",":","`","256","`",")","num_spherical","(","int",",","optional",")",":","Number","of","spherical","harmonics",".","(","default",":",":","obj",":","`","7","`",")","num_radial","(","int",",","optional",")",":","Number","of","radial","basis","functions",".","(","default",":",":","obj",":","`","6","`",")","envelop_exponent","(","int",",","optional",")",":","Shape","of","the","smooth","cutoff",".","(","default",":",":","obj",":","`","5","`",")","num_before_skip","(","int",",","optional",")",":","Number","of","residual","layers","in","the","interaction","blocks","before","the","skip","connection",".","(","default",":",":","obj",":","`","1","`",")","num_after_skip","(","int",",","optional",")",":","Number","of","residual","layers","in","the","interaction","blocks","before","the","skip","connection",".","(","default",":",":","obj",":","`","2","`",")","num_output_layers","(","int",",","optional",")",":","Number","of","linear","layers","for","the","output","blocks",".","(","default",":",":","obj",":","`","3","`",")","act",":","(","function",",","optional",")",":","The","activation","funtion",".","(","default",":",":","obj",":","`","swish","`",")","output_init",":","(","str",",","optional",")",":","The","initialization","fot","the","output",".","It","could","be",":","obj",":","`","GlorotOrthogonal","`","and",":","obj",":","`","zeros","`",".","(","default",":",":","obj",":","`","GlorotOrthogonal","`",")"]
code: def forward(self, batch_data):
      z, pos, batch = batch_data.z, batch_data.pos, batch_data.batch
      if self.energy_and_force:
          pos.requires_grad_()
      edge_index = radius_graph(pos, r=self.cutoff, batch=batch)
      num_nodes = z.size(0)
      dist, angle, i, j, idx_kj, idx_ji = xyz_to_dat(
          pos, edge_index, num_nodes, use_torsion=False
      )
  
      emb = self.emb(dist, angle, idx_kj)
  
      # Initialize edge, node, graph features
      e = self.init_e(z, emb, i, j)
      v = self.init_v(e, i)
      u = self.init_u(
          torch.zeros_like(scatter(v, batch, dim=0)), v, batch
      )  # scatter(v, batch, dim=0)
  
      for update_e, update_v, update_u in zip(
          self.update_es, self.update_vs, self.update_us
      ):
          e = update_e(e, emb, idx_kj, idx_ji)
          v = update_v(e, i)
          u = update_u(u, v, batch)  # u += scatter(v, batch, dim=0)
  
      return u
  
imports: import torch
  from torch import nn
  from torch.nn import Linear, Embedding
  from torch_geometric.nn.inits import glorot_orthogonal
  from torch_geometric.nn import radius_graph
  from torch_scatter import scatter
  from math import sqrt
  from ...utils import xyz_to_dat
  from .features import dist_emb, angle_emb
label_desc: Use image node_id 3 for calling the DimeNetPP obj's underlying member method code with example usage: obj.forward(batch_data) and returns: u