code: def __init__(self, vocab, hidden_size, latent_size, depth):
      super(DGLJTNNVAE, self).__init__()
      self.vocab = vocab
      self.hidden_size = hidden_size
      self.latent_size = latent_size
      self.depth = depth
  
      self.embedding = nn.Embedding(vocab.size(), hidden_size)
      self.mpn = DGLMPN(hidden_size, depth)
      self.jtnn = DGLJTNNEncoder(vocab, hidden_size, self.embedding)
      self.decoder = DGLJTNNDecoder(
          vocab, hidden_size, latent_size // 2, self.embedding
      )
      self.jtmpn = DGLJTMPN(hidden_size, depth)
  
      self.T_mean = nn.Linear(hidden_size, latent_size // 2)
      self.T_var = nn.Linear(hidden_size, latent_size // 2)
      self.G_mean = nn.Linear(hidden_size, latent_size // 2)
      self.G_var = nn.Linear(hidden_size, latent_size // 2)
  
      self.n_nodes_total = 0
      self.n_passes = 0
      self.n_edges_total = 0
      self.n_tree_nodes_total = 0
  
imports: import copy
  import rdkit.Chem
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import batch, unbatch
  from .chemutils import attach_mols_nx, copy_edit_mol, decode_stereo, enum_assemble_nx, set_atommap
  from .jtmpn import DGLJTMPN, mol2dgl_single
  from .jtnn_dec import DGLJTNNDecoder
  from .jtnn_enc import DGLJTNNEncoder
  from .mpn import DGLMPN, mol2dgl_single
  from .nnutils import cuda
label_desc: Use image node_id 1 to create a new DGLJTNNVAE object from inherited base classes: nn with example: obj = DGLJTNNVAE(vocab, hidden_size, latent_size, depth)