class_docstr: The LINKX model from the `"Large Scale Learning on Non-Homophilous
  Graphs: New Benchmarks and Strong Simple Methods"
  <https://arxiv.org/abs/2110.14446>`_ paper.
  
  .. math::
      \mathbf{H}_{\mathbf{A}} &= \textrm{MLP}_{\mathbf{A}}(\mathbf{A})
  
      \mathbf{H}_{\mathbf{X}} &= \textrm{MLP}_{\mathbf{X}}(\mathbf{X})
  
      \mathbf{Y} &= \textrm{MLP}_{f} \left( \sigma \left( \mathbf{W}
      [\mathbf{H}_{\mathbf{A}}, \mathbf{H}_{\mathbf{X}}] +
      \mathbf{H}_{\mathbf{A}} + \mathbf{H}_{\mathbf{X}} \right) \right)
  
  .. note::
  
      For an example of using LINKX, see `examples/linkx.py <https://
      github.com/pyg-team/pytorch_geometric/blob/master/examples/linkx.py>`_.
  
  Args:
      num_nodes (int): The number of nodes in the graph.
      in_channels (int): Size of each input sample, or :obj:`-1` to derive
          the size from the first input(s) to the forward method.
      hidden_channels (int): Size of each hidden sample.
      out_channels (int): Size of each output sample.
      num_layers (int): Number of layers of :math:`\textrm{MLP}_{f}`.
      num_edge_layers (int, optional): Number of layers of
          :math:`\textrm{MLP}_{\mathbf{A}}`. (default: :obj:`1`)
      num_node_layers (int, optional): Number of layers of
          :math:`\textrm{MLP}_{\mathbf{X}}`. (default: :obj:`1`)
      dropout (float, optional): Dropout probability of each hidden
          embedding. (default: :obj:`0.0`)
code: def __init__(
      self,
      num_nodes: int,
      in_channels: int,
      hidden_channels: int,
      out_channels: int,
      num_layers: int,
      num_edge_layers: int = 1,
      num_node_layers: int = 1,
      dropout: float = 0.0,
  ):
      super().__init__()
  
      self.num_nodes = num_nodes
      self.in_channels = in_channels
      self.out_channels = out_channels
      self.num_edge_layers = num_edge_layers
  
      self.edge_lin = SparseLinear(num_nodes, hidden_channels)
  
      if self.num_edge_layers > 1:
          self.edge_norm = BatchNorm1d(hidden_channels)
          channels = [hidden_channels] * num_edge_layers
          self.edge_mlp = MLP(channels, dropout=0.0, act_first=True)
      else:
          self.edge_norm = None
          self.edge_mlp = None
  
      channels = [in_channels] + [hidden_channels] * num_node_layers
      self.node_mlp = MLP(channels, dropout=0.0, act_first=True)
  
      self.cat_lin1 = torch.nn.Linear(hidden_channels, hidden_channels)
      self.cat_lin2 = torch.nn.Linear(hidden_channels, hidden_channels)
  
      channels = [hidden_channels] * num_layers + [out_channels]
      self.final_mlp = MLP(channels, dropout=dropout, act_first=True)
  
      self.reset_parameters()
  
imports: import math
  import torch
  from torch import Tensor
  from torch.nn import BatchNorm1d, Parameter
  from torch_geometric.nn import inits
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.nn.models import MLP
  from torch_geometric.typing import Adj, OptTensor
  from torch_geometric.utils import spmm
label_desc: Use image node_id 1 to create a new LINKX object from inherited base classes: torch.nn with example: obj = LINKX(num_nodes, in_channels, hidden_channels, out_channels, num_layers, num_edge_layers, num_node_layers, dropout)