code: def __init__(self, vocab, hidden_size, latent_size, depthT, depthG):
      super().__init__()
      self.vocab = vocab
      self.hidden_size = hidden_size
      self.latent_size = latent_size = (
          latent_size // 2
      )  # Tree and Mol has two vectors
  
      self.jtnn = JTNNEncoder(
          hidden_size, depthT, nn.Embedding(vocab.size(), hidden_size)
      )
      self.decoder = JTNNDecoder(
          vocab,
          hidden_size,
          latent_size,
          nn.Embedding(vocab.size(), hidden_size),
      )
  
      self.jtmpn = JTMPN(hidden_size, depthG)
      self.mpn = MPN(hidden_size, depthG)
  
      self.A_assm = nn.Linear(latent_size, hidden_size, bias=False)
      self.assm_loss = nn.CrossEntropyLoss(size_average=False)
  
      self.T_mean = nn.Linear(hidden_size, latent_size)
      self.T_var = nn.Linear(hidden_size, latent_size)
      self.G_mean = nn.Linear(hidden_size, latent_size)
      self.G_var = nn.Linear(hidden_size, latent_size)
  
imports: import copy
  import math
  import torch
  import torch.nn
  import torch.nn.functional
  import rdkit
  import rdkit.Chem
  from .mol_tree import MolTree
  from .nnutils import create_var
  from .jtnn_enc import JTNNEncoder
  from .jtnn_dec import JTNNDecoder
  from .mpn import MPN
  from .jtmpn import JTMPN
  from .datautils import tensorize
  from .chemutils import enum_assemble, set_atommap, copy_edit_mol, attach_mols
label_desc: Use image node_id 1 to create a new JTNNVAE object from inherited base classes: nn with example: obj = JTNNVAE(vocab, hidden_size, latent_size, depthT, depthG)