class_docstr: DimeNet model.
  
  Parameters
  ----------
  emb_size
      Embedding size used throughout the model
  num_blocks
      Number of building blocks to be stacked
  num_bilinear
      Third dimension of the bilinear layer tensor
  num_spherical
      Number of spherical harmonics
  num_radial
      Number of radial basis functions
  cutoff
      Cutoff distance for interatomic interactions
  envelope_exponent
      Shape of the smooth cutoff
  num_before_skip
      Number of residual layers in interaction block before skip connection
  num_after_skip
      Number of residual layers in interaction block after skip connection
  num_dense_output
      Number of dense layers for the output blocks
  num_targets
      Number of targets to predict
  activation
      Activation function
  output_init
      Initial function in output block
class_docstr_tok: ["DimeNet","model",".","Parameters","--","--","--","--","--","emb_size","Embedding","size","used","throughout","the","model","num_blocks","Number","of","building","blocks","to","be","stacked","num_bilinear","Third","dimension","of","the","bilinear","layer","tensor","num_spherical","Number","of","spherical","harmonics","num_radial","Number","of","radial","basis","functions","cutoff","Cutoff","distance","for","interatomic","interactions","envelope_exponent","Shape","of","the","smooth","cutoff","num_before_skip","Number","of","residual","layers","in","interaction","block","before","skip","connection","num_after_skip","Number","of","residual","layers","in","interaction","block","after","skip","connection","num_dense_output","Number","of","dense","layers","for","the","output","blocks","num_targets","Number","of","targets","to","predict","activation","Activation","function","output_init","Initial","function","in","output","block"]
code: def edge_init(self, edges):
      # Calculate angles k -> j -> i
      R1, R2 = edges.src["o"], edges.dst["o"]
      x = torch.sum(R1 * R2, dim=-1)
      y = torch.cross(R1, R2)
      y = torch.norm(y, dim=-1)
      angle = torch.atan2(y, x)
      # Transform via angles
      cbf = [f(angle) for f in self.sbf_layer.get_sph_funcs()]
      cbf = torch.stack(cbf, dim=1)  # [None, 7]
      cbf = cbf.repeat_interleave(self.num_radial, dim=1)  # [None, 42]
      sbf = edges.src["rbf_env"] * cbf  # [None, 42]
      return {"sbf": sbf}
  
file_path: dgl/examples/pytorch/dimenet/modules/dimenet.py
imports: import torch
  import torch.nn
  from modules.activations import swish
  from modules.bessel_basis_layer import BesselBasisLayer
  from modules.embedding_block import EmbeddingBlock
  from modules.interaction_block import InteractionBlock
  from modules.output_block import OutputBlock
  from modules.spherical_basis_layer import SphericalBasisLayer
label_desc: Use image node_id 2 for calling the DimeNet obj's underlying member method code with example usage: obj.edge_init(edges) and returns: dict