class_docstr: Net3DLayer is a single layer of a 3D graph neural network based on the 3D Infomax architecture [1].
  
  This class expects a DGL graph with node features stored under the name 'feat' and edge features stored under the name 'd' (representing 3D distances). The edge features are updated by the message network and the node features are updated by the update network.
  
  Parameters
  ----------
  edge_dim : int
      The dimension of the edge features.
  hidden_dim : int
      The dimension of the hidden layers.
  reduce_func : str
      The reduce function to use for aggregating messages. Can be either 'sum' or 'mean'.
  batch_norm : bool, optional (default=False)
      Whether to use batch normalization.
  batch_norm_momentum : float, optional (default=0.1)
      The momentum for the batch normalization layers.
  dropout : float, optional (default=0.0)
      The dropout rate for the layers.
  mid_activation : str, optional (default='SiLU')
      The activation function to use in the network.
  message_net_layers : int, optional (default=2)
      The number of message network layers.
  update_net_layers : int, optional (default=2)
      The number of update network layers.
  
  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
  --------
  >>> net3d_layer = Net3DLayer(edge_dim=3, hidden_dim=3)
  >>> graph = dgl.graph(([0, 1], [1, 2]))
  >>> graph.ndata['feat'] = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
  >>> graph.edata['d'] = torch.tensor([[0.5, 0.6, 0.7], [0.8, 0.9, 1.0]])
  >>> output = net3d_layer(graph)
code: def __init__(
      self,
      edge_dim: int,
      hidden_dim: int,
      reduce_func: str = "sum",
      batch_norm: bool = False,
      batch_norm_momentum: float = 0.1,
      dropout: float = 0.0,
      message_net_layers: int = 2,
      update_net_layers: int = 2,
  ):
      super(Net3DLayer, self).__init__()
  
      self.message_network = nn.Sequential(
          MultilayerPerceptron(
              d_input=hidden_dim * 2 + edge_dim,
              d_output=hidden_dim,
              d_hidden=(hidden_dim,) * (message_net_layers - 1),
              batch_norm=batch_norm,
              batch_norm_momentum=batch_norm_momentum,
              dropout=dropout,
          ),
          torch.nn.SiLU(),
      )
      if reduce_func == "sum":
          self.reduce_func = fn.sum
      elif reduce_func == "mean":
          self.reduce_func = fn.mean
      else:
          raise ValueError(
              "reduce function not supported: ", reduce_func
          )
  
      self.update_network = MultilayerPerceptron(
          d_input=hidden_dim,
          d_hidden=(hidden_dim,) * (update_net_layers - 1),
          d_output=hidden_dim,
          batch_norm=True,
          batch_norm_momentum=batch_norm_momentum,
      )
  
      self.soft_edge_network = nn.Linear(hidden_dim, 1)
  
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 1 to create a new Net3DLayer object from inherited base classes: nn with example: obj = Net3DLayer(edge_dim, hidden_dim, reduce_func, batch_norm, batch_norm_momentum, dropout, message_net_layers, update_net_layers)