class_docstr: Principal Neighbourhood Aggregation Layer (PNA) from [1].
  
  Parameters
  ----------
  in_dim : int
      Input dimension of the node features.
  out_dim : int
      Output dimension of the node features.
  in_dim_edges : int
      Input dimension of the edge features.
  aggregators : List[str]
      List of aggregator functions to use. Options are "mean", "sum", "max", "min", "std", "var", "moment3", "moment4", "moment5".
  scalers : List[str]
      List of scaler functions to use. Options are "identity", "amplification", "attenuation".
  activation : Union[Callable, str], optional, default="relu"
      Activation function to use.
  last_activation : Union[Callable, str], optional, default="none"
      Last activation function to use.
  dropout : float, optional, default=0.0
      Dropout rate.
  residual : bool, optional, default=True
      Whether to use residual connections.
  pairwise_distances : bool, optional, default=False
      Whether to use pairwise distances.
  batch_norm : bool, optional, default=True
      Whether to use batch normalization.
  batch_norm_momentum : float, optional, default=0.1
      Momentum for the batch normalization layers.
  avg_d : Dict[str, float], optional, default={"log": 1.0}
      Dictionary containing the average degree of the graph.
  posttrans_layers : int, optional, default=2
      Number of post-transformation layers.
  pretrans_layers : int, optional, 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 dgl
  >>> import numpy as np
  >>> import torch
  >>> from deepchem.models.torch_models.pna_gnn import PNALayer
  >>> in_dim = 32
  >>> out_dim = 64
  >>> in_dim_edges = 16
  >>> aggregators = ["mean", "max"]
  >>> scalers = ["identity", "amplification", "attenuation"]
  >>> pna_layer = PNALayer(in_dim=in_dim,
  ...                      out_dim=out_dim,
  ...                      in_dim_edges=in_dim_edges,
  ...                      aggregators=aggregators,
  ...                      scalers=scalers)
  >>> num_nodes = 10
  >>> num_edges = 20
  >>> node_features = torch.randn(num_nodes, in_dim)
  >>> edge_features = torch.randn(num_edges, in_dim_edges)
  >>> g = dgl.graph((np.random.randint(0, num_nodes, num_edges),
  ...                np.random.randint(0, num_nodes, num_edges)))
  >>> g.ndata['feat'] = node_features
  >>> g.edata['feat'] = edge_features
  >>> g.ndata['feat'] = pna_layer(g)
class_docstr_tok: ["Principal","Neighbourhood","Aggregation","Layer","(","PNA",")","from","[","1","]",".","Parameters","--","--","--","--","--","in_dim",":","int","Input","dimension","of","the","node","features",".","out_dim",":","int","Output","dimension","of","the","node","features",".","in_dim_edges",":","int","Input","dimension","of","the","edge","features",".","aggregators",":","List","[","str","]","List","of","aggregator","functions","to","use",".","Options","are","``","mean","''",",","``","sum","''",",","``","max","''",",","``","min","''",",","``","std","''",",","``","var","''",",","``","moment3","''",",","``","moment4","''",",","``","moment5","''",".","scalers",":","List","[","str","]","List","of","scaler","functions","to","use",".","Options","are","``","identity","''",",","``","amplification","''",",","``","attenuation","''",".","activation",":","Union","[","Callable",",","str","]",",","optional",",","default=","''","relu","''","Activation","function","to","use",".","last_activation",":","Union","[","Callable",",","str","]",",","optional",",","default=","''","none","''","Last","activation","function","to","use",".","dropout",":","float",",","optional",",","default=0.0","Dropout","rate",".","residual",":","bool",",","optional",",","default=True","Whether","to","use","residual","connections",".","pairwise_distances",":","bool",",","optional",",","default=False","Whether","to","use","pairwise","distances",".","batch_norm",":","bool",",","optional",",","default=True","Whether","to","use","batch","normalization",".","batch_norm_momentum",":","float",",","optional",",","default=0.1","Momentum","for","the","batch","normalization","layers",".","avg_d",":","Dict","[","str",",","float","]",",","optional",",","default=","{","``","log","''",":","1.0","}","Dictionary","containing","the","average","degree","of","the","graph",".","posttrans_layers",":","int",",","optional",",","default=2","Number","of","post-transformation","layers",".","pretrans_layers",":","int",",","optional",",","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","dgl",">",">",">","import","numpy","as","np",">",">",">","import","torch",">",">",">","from","deepchem.models.torch_models.pna_gnn","import","PNALayer",">",">",">","in_dim","=","32",">",">",">","out_dim","=","64",">",">",">","in_dim_edges","=","16",">",">",">","aggregators","=","[","``","mean","''",",","``","max","''","]",">",">",">","scalers","=","[","``","identity","''",",","``","amplification","''",",","``","attenuation","''","]",">",">",">","pna_layer","=","PNALayer","(","in_dim=in_dim",",","...","out_dim=out_dim",",","...","in_dim_edges=in_dim_edges",",","...","aggregators=aggregators",",","...","scalers=scalers",")",">",">",">","num_nodes","=","10",">",">",">","num_edges","=","20",">",">",">","node_features","=","torch.randn","(","num_nodes",",","in_dim",")",">",">",">","edge_features","=","torch.randn","(","num_edges",",","in_dim_edges",")",">",">",">","g","=","dgl.graph","(","(","np.random.randint","(","0",",","num_nodes",",","num_edges",")",",","...","np.random.randint","(","0",",","num_nodes",",","num_edges",")",")",")",">",">",">","g.ndata","[","'feat","'","]","=","node_features",">",">",">","g.edata","[","'feat","'","]","=","edge_features",">",">",">","g.ndata","[","'feat","'","]","=","pna_layer","(","g",")"]
docstr: The reduce function to aggregate the messages.
  Apply the aggregators and scalers, and concatenate the results.
  
  Parameters
  ----------
  nodes : dgl.NodeBatch
      Batch of nodes.
  
  Returns
  -------
  Dict[str, torch.Tensor]
      Dictionary containing the aggregated node features.
docstr_tok: ["The","reduce","function","to","aggregate","the","messages",".","Apply","the","aggregators","and","scalers",",","and","concatenate","the","results",".","Parameters","--","--","--","--","--","nodes",":","dgl.NodeBatch","Batch","of","nodes",".","Returns","--","--","--","-","Dict","[","str",",","torch.Tensor","]","Dictionary","containing","the","aggregated","node","features","."]
code: def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
      """
      The reduce function to aggregate the messages.
      Apply the aggregators and scalers, and concatenate the results.
  
      Parameters
      ----------
      nodes : dgl.NodeBatch
          Batch of nodes.
  
      Returns
      -------
      Dict[str, torch.Tensor]
          Dictionary containing the aggregated node features.
      """
      h_in = nodes.data["feat"]
      h = nodes.mailbox["e"]
      D = h.shape[-2]
  
      h_to_cat = [
          aggr(h=h, h_in=h_in)  # type: ignore
          for aggr in self.aggregators
      ]
      h = torch.cat(h_to_cat, dim=-1)
  
      if len(self.scalers) > 1:
          h = torch.cat(
              [
                  scale(h, D=D, avg_d=self.avg_d)  # type: ignore
                  for scale in self.scalers
              ],
              dim=-1,
          )
  
      return {"feat": h}
  
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 4 for calling the PNALayer obj's underlying member method code with example usage: obj.reduce_func(nodes) and returns: dict