class_docstr: Message passing neural network for graph representation learning [1]_.
  
  Parameters
  ----------
  hidden_dim : int
      Hidden dimension size.
  target_dim : int
      Dimensionality of the output, for example for binary classification target_dim = 1.
  aggregators : List[str]
      Type of message passing functions. Options are 'mean','sum','max','min','std','var','moment3','moment4','moment5'.
  scalers : List[str]
      Type of normalization layers in the message passing network. Options are 'identity','amplification','attenuation'.
  readout_aggregators : List[str]
      Type of aggregators in the readout network.
  readout_hidden_dim : int, default None
     The dimension of the hidden layer in the readout network. If not provided, the readout has the same dimensionality of the final layer of the PNA layer, which is the hidden dimension size.
  readout_layers : int, default 1
      The number of linear layers in the readout network.
  residual : bool, default True
      Whether to use residual connections.
  pairwise_distances : bool, default False
      Whether to use pairwise distances.
  activation : Union[Callable, str]
      Activation function to use.
  batch_norm : bool, default True
      Whether to use batch normalization in the layers before the aggregator..
  batch_norm_momentum : float, default 0.1
      Momentum for the batch normalization layers.
  propagation_depth : int, default
      Number of propagation layers.
  dropout : float, default 0.0
      Dropout probability in the message passing layers.
  posttrans_layers : int, default 1
      Number of post-transformation layers.
  pretrans_layers : int, default 1
      Number of pre-transformation layers.
  
  References
  ----------
  .. [1] Corso, G., Cavalleri, L., Beaini, D., Liò, P. & Veličković, P. Principal Neighbourhood Aggregation for Graph Nets. Preprint at https://doi.org/10.48550/arXiv.2004.05718 (2020).
  
  Examples
  --------
  >>> import numpy as np
  >>> from deepchem.feat.graph_data import BatchGraphData
  >>> from deepchem.models.torch_models.pna_gnn import PNA
  >>> from deepchem.feat.molecule_featurizers.conformer_featurizer import RDKitConformerFeaturizer
  >>> smiles = ["C1=CC=CN=C1", "C1CCC1"]
  >>> featurizer = RDKitConformerFeaturizer(num_conformers=2)
  >>> data = featurizer.featurize(smiles)
  >>> features = BatchGraphData(np.concatenate(data))
  >>> features = features.to_dgl_graph()
  >>> target_dim = 1
  >>> model = PNA(hidden_dim=16, target_dim=target_dim)
  >>> output = model(features)
  >>> print(output.shape)
  torch.Size([1, 1])
class_docstr_tok: ["Message","passing","neural","network","for","graph","representation","learning","[","1","]","_.","Parameters","--","--","--","--","--","hidden_dim",":","int","Hidden","dimension","size",".","target_dim",":","int","Dimensionality","of","the","output",",","for","example","for","binary","classification","target_dim","=","1.","aggregators",":","List","[","str","]","Type","of","message","passing","functions",".","Options","are","'mean","'",",","'sum","'",",","'max","'",",","'min","'",",","'std","'",",","'var","'",",","'moment3","'",",","'moment4","'",",","'moment5","'",".","scalers",":","List","[","str","]","Type","of","normalization","layers","in","the","message","passing","network",".","Options","are","'identity","'",",","'amplification","'",",","'attenuation","'",".","readout_aggregators",":","List","[","str","]","Type","of","aggregators","in","the","readout","network",".","readout_hidden_dim",":","int",",","default","None","The","dimension","of","the","hidden","layer","in","the","readout","network",".","If","not","provided",",","the","readout","has","the","same","dimensionality","of","the","final","layer","of","the","PNA","layer",",","which","is","the","hidden","dimension","size",".","readout_layers",":","int",",","default","1","The","number","of","linear","layers","in","the","readout","network",".","residual",":","bool",",","default","True","Whether","to","use","residual","connections",".","pairwise_distances",":","bool",",","default","False","Whether","to","use","pairwise","distances",".","activation",":","Union","[","Callable",",","str","]","Activation","function","to","use",".","batch_norm",":","bool",",","default","True","Whether","to","use","batch","normalization","in","the","layers","before","the","aggregator","..","batch_norm_momentum",":","float",",","default","0.1","Momentum","for","the","batch","normalization","layers",".","propagation_depth",":","int",",","default","Number","of","propagation","layers",".","dropout",":","float",",","default","0.0","Dropout","probability","in","the","message","passing","layers",".","posttrans_layers",":","int",",","default","1","Number","of","post-transformation","layers",".","pretrans_layers",":","int",",","default","1","Number","of","pre-transformation","layers",".","References","--","--","--","--","--","..","[","1","]","Corso",",","G.",",","Cavalleri",",","L.",",","Beaini",",","D.",",","Li\u00f2",",","P.","&","Veli\u010dkovi\u0107",",","P.","Principal","Neighbourhood","Aggregation","for","Graph","Nets",".","Preprint","at","https",":","\/\/doi.org\/10.48550\/arXiv.2004.05718","(","2020",")",".","Examples","--","--","--","--",">",">",">","import","numpy","as","np",">",">",">","from","deepchem.feat.graph_data","import","BatchGraphData",">",">",">","from","deepchem.models.torch_models.pna_gnn","import","PNA",">",">",">","from","deepchem.feat.molecule_featurizers.conformer_featurizer","import","RDKitConformerFeaturizer",">",">",">","smiles","=","[","``","C1=CC=CN=C1","''",",","``","C1CCC1","''","]",">",">",">","featurizer","=","RDKitConformerFeaturizer","(","num_conformers=2",")",">",">",">","data","=","featurizer.featurize","(","smiles",")",">",">",">","features","=","BatchGraphData","(","np.concatenate","(","data",")",")",">",">",">","features","=","features.to_dgl_graph","(",")",">",">",">","target_dim","=","1",">",">",">","model","=","PNA","(","hidden_dim=16",",","target_dim=target_dim",")",">",">",">","output","=","model","(","features",")",">",">",">","print","(","output.shape",")","torch.Size","(","[","1",",","1","]",")"]
code: def forward(self, graph: dgl.DGLGraph):
      graph = self.node_gnn(graph)
      readouts_to_cat = [
          dgl.readout_nodes(graph, "feat", op=aggr)
          for aggr in self.readout_aggregators
      ]
      readout = torch.cat(readouts_to_cat, dim=-1)
      outputs = self.output(readout)
      if self.task == "classification":
          if self.n_tasks == 1:
              softmax_dim = 1
              logits = outputs.view(-1, self.n_classes)
          else:
              softmax_dim = 2
              logits = outputs.view(-1, self.n_tasks, self.n_classes)
          proba = F.softmax(logits, dim=softmax_dim)
          # print (logits, proba)
          return proba, logits
      else:
          return outputs
  
imports: from functools import partial
  from math import sqrt
  from typing import Callable, Dict, List, Union
  import dgl
  import torch
  from torch import nn
  import torch.nn.functional
  from deepchem.feat.molecule_featurizers.conformer_featurizer import full_atom_feature_dims, full_bond_feature_dims
  from deepchem.models.torch_models.layers import MultilayerPerceptron
  from deepchem.utils.graph_utils import aggregate_max, aggregate_mean, aggregate_min, aggregate_moment, aggregate_std, aggregate_sum, aggregate_var, scale_amplification, scale_attenuation, scale_identity
label_desc: Use image node_id 2 for calling the PNA obj's underlying member method code with example usage: obj.forward(graph) and returns: proba,  logits, outputs