class_docstr: InfoMax3DModular is a modular torch model that uses a 2D PNA model and a 3D Net3D model to maximize the mutual information between their representations. The 2D model can then be used for downstream tasks without the need for 3D coordinates. This is based off the work in [1].
  
  This class expects data in featurized by the RDKitConformerFeaturizer. This featurizer produces features of the type Array[Array[List[GraphData]]].
  The outermost array is the dataset array, the second array is the molecule, the list contains the conformers for that molecule and the GraphData object is the featurized graph for that conformer with node_pos_features holding the 3D coordinates.
  If you are not using RDKitConformerFeaturizer, your input data features should look like this: Dataset[Molecule[Conformers[GraphData]]].
  
  For pretraining, the original paper used a learning rate of 8e-5 with a batch size of 500.
  For finetuning on quantum mechanical datasets, a learning rate of 7e-5 with a batch size of 128
  was used. For finetuning on non-quantum mechanical datasets, a learning rate of 1e-3 with a
  batch size of 32 was used in the original implementation.
  
  Parameters
  ----------
  task : Literal['pretrain', 'regression', 'classification']
      The task of the model
  hidden_dim : int, optional, default = 64
      The dimension of the hidden layers.
  target_dim : int, optional, default = 10
      The dimension of the output layer.
  aggregators : List[str]
      A list of aggregator functions for the PNA model. Options are 'mean', 'sum', 'min', 'max', 'std', 'var', 'moment3', 'moment4', 'moment5'.
  readout_aggregators : List[str]
      A list of aggregator functions for the readout layer. Options are 'sum', 'max', 'min', 'mean'.
  scalers : List[str]
      A list of scaler functions for the PNA model. Options are 'identity', 'amplification', 'attenuation'.
  residual : bool, optional (default=True)
      Whether to use residual connections in the PNA model.
  node_wise_output_layers : int, optional (default=2)
      The number of output layers for each node in the Net3D model.
  pairwise_distances : bool, optional (default=False)
      Whether to use pairwise distances in the PNA model.
  activation : Union[Callable, str], optional (default="relu")
      The activation function to use in the PNA model.
  reduce_func : str, optional (default='sum')
      The reduce function to use for aggregating messages in the Net3D model.
  batch_norm : bool, optional (default=True)
      Whether to use batch normalization in the PNA model.
  batch_norm_momentum : float, optional (default=0.1)
      The momentum for the batch normalization layers.
  propagation_depth : int, optional (default=5)
      The number of propagation layers in the PNA and Net3D models.
  dropout : float, optional (default=0.0)
      The dropout rate for the layers in the PNA and Net3D models.
  readout_layers : int, optional (default=2)
      The number of readout layers in the PNA and Net3D models.
  readout_hidden_dim : int, optional (default=None)
      The dimension of the hidden layers in the readout network.
  fourier_encodings : int, optional (default=4)
      The number of Fourier encodings to use in the Net3D model.
  update_net_layers : int, optional (default=2)
      The number of update network layers in the Net3D model.
  message_net_layers : int, optional (default=2)
      The number of message network layers in the Net3D model.
  use_node_features : bool, optional (default=False)
      Whether to use node features as input in the Net3D model.
  posttrans_layers : int, optional (default=1)
      The number of post-transformation layers in the PNA model.
  pretrans_layers : int, optional (default=1)
      The number of pre-transformation layers in the PNA model.
  kwargs : dict
      Additional keyword arguments.
  
  References
  ----------
  .. [1] Stärk, H. et al. 3D Infomax improves GNNs for Molecular Property Prediction. Preprint at https://doi.org/10.48550/arXiv.2110.04126 (2022).
  
  Examples
  --------
  >>> from deepchem.feat.graph_data import BatchGraphData
  >>> from deepchem.feat.molecule_featurizers.conformer_featurizer import RDKitConformerFeaturizer
  >>> from deepchem.models.torch_models.gnn3d import InfoMax3DModular
  >>> import numpy as np
  >>> import deepchem as dc
  >>> from deepchem.data.datasets import NumpyDataset
  >>> smiles = ["C[C@H](F)Cl", "C[C@@H](F)Cl"]
  >>> featurizer = RDKitConformerFeaturizer()
  >>> data = featurizer.featurize(smiles)
  >>> dataset = NumpyDataset(X=data)
  >>> model = InfoMax3DModular(task='pretraining',
  ...                          hidden_dim=64,
  ...                          target_dim=10,
  ...                          aggregators=['max'],
  ...                          readout_aggregators=['mean'],
  ...                          scalers=['identity'])
  >>> loss = model.fit(dataset, nb_epoch=1)
docstr: Build the InfoMax3DModular model. This is the 2D network which is meant to be used for inference.
  
  Returns
  -------
  PNA
      The 2D PNA model component.
docstr_tok: ["Build","the","InfoMax3DModular","model",".","This","is","the","2D","network","which","is","meant","to","be","used","for","inference",".","Returns","--","--","--","-","PNA","The","2D","PNA","model","component","."]
code: def build_model(self):
      """
      Build the InfoMax3DModular model. This is the 2D network which is meant to be used for inference.
  
      Returns
      -------
      PNA
          The 2D PNA model component.
      """
      # FIXME For pretraining task, both model2d and model3d but the super class
      # can't handle two models for contrastive learning, hence we pass only model2d
      return self.components["model2d"]
  
imports: import itertools
  from typing import Callable, List, Union, Literal, Optional
  import dgl
  import dgl.function
  import torch
  from torch import nn
  from torch.nn import functional
  from deepchem.models.losses import NTXentMultiplePositives
  from deepchem.models.torch_models import ModularTorchModel
  from deepchem.models.torch_models.layers import MultilayerPerceptron
  from deepchem.models.torch_models.pna_gnn import PNA, AtomEncoder
  from deepchem.utils.graph_utils import fourier_encode_dist
label_desc: Use image node_id 3 for calling the InfoMax3DModular obj's underlying member method code with example usage: obj.build_model() and returns: self