code: def summary(
      model: torch.nn.Module,
      *args,
      max_depth: int = 3,
      leaf_module: Optional[
          Union[Module, List[Module]]
      ] = "MessagePassing",
      **kwargs,
  ) -> str:
      r"""Summarizes a given :class:`torch.nn.Module`.
      The summarized information includes (1) layer names, (2) input and output
      shapes, and (3) the number of parameters.
  
      .. code-block:: python
  
          import torch
          from torch_geometric.nn import GCN, summary
  
          model = GCN(128, 64, num_layers=2, out_channels=32)
          x = torch.randn(100, 128)
          edge_index = torch.randint(100, size=(2, 20))
  
          print(summary(model, x, edge_index))
  
      .. code-block::
  
          +---------------------+---------------------+--------------+--------+
          | Layer               | Input Shape         | Output Shape | #Param |
          |---------------------+---------------------+--------------+--------|
          | GCN                 | [100, 128], [2, 20] | [100, 32]    | 10,336 |
          | ├─(act)ReLU         | [100, 64]           | [100, 64]    | --     |
          | ├─(convs)ModuleList | --                  | --           | 10,336 |
          | │    └─(0)GCNConv   | [100, 128], [2, 20] | [100, 64]    | 8,256  |
          | │    └─(1)GCNConv   | [100, 64], [2, 20]  | [100, 32]    | 2,080  |
          +---------------------+---------------------+--------------+--------+
  
      Args:
          model (torch.nn.Module): The model to summarize.
          *args: The arguments of the :obj:`model`.
          max_depth (int, optional): The depth of nested layers to display.
              Any layers deeper than this depth will not be displayed in the
              summary. (default: :obj:`3`)
          leaf_module (torch.nn.Module or [torch.nn.Module], optional): The
              modules to be treated as leaf modules, whose submodules are
              excluded from the summary.
              (default: :class:`~torch_geometric.nn.conv.MessagePassing`)
          **kwargs: Additional arguments of the :obj:`model`.
      """
      # NOTE This is just for the doc-string to render nicely:
      if leaf_module == "MessagePassing":
          leaf_module = MessagePassing
  
      def register_hook(info):
          def hook(module, inputs, output):
              info["input_shape"].append(get_shape(inputs))
              info["output_shape"].append(get_shape(output))
  
          return hook
  
      hooks = {}
      depth = 0
      stack = [(model.__class__.__name__, model, depth)]
  
      info_list = []
      input_shape = defaultdict(list)
      output_shape = defaultdict(list)
      while stack:
          name, module, depth = stack.pop()
          module_id = id(module)
  
          if name.startswith("(_"):  # Do not summarize private modules.
              continue
  
          if module_id in hooks:  # Avoid duplicated hooks.
              hooks[module_id].remove()
  
          info = {}
          info["name"] = name
          info["input_shape"] = input_shape[module_id]
          info["output_shape"] = output_shape[module_id]
          info["depth"] = depth
          if any(
              [
                  is_uninitialized_parameter(p)
                  for p in module.parameters()
              ]
          ):
              info["#param"] = "-1"
          else:
              num_params = sum(p.numel() for p in module.parameters())
              info["#param"] = (
                  f"{num_params:,}" if num_params > 0 else "--"
              )
          info_list.append(info)
  
          if not isinstance(module, ScriptModule):
              hooks[module_id] = module.register_forward_hook(
                  register_hook(info)
              )
  
          if depth >= max_depth:
              continue
  
          if leaf_module is not None and isinstance(
              module, leaf_module
          ):
              continue
  
          module_items = reversed(module._modules.items())
          stack += [
              (f"({name}){mod.__class__.__name__}", mod, depth + 1)
              for name, mod in module_items
              if mod is not None
          ]
  
      training = model.training
      model.eval()
  
      with torch.no_grad():
          model(*args, **kwargs)
  
      model.train(training)
  
      for h in hooks.values():  # Remove hooks.
          h.remove()
  
      info_list = postprocess(info_list)
      return make_table(info_list, max_depth=max_depth)
  
file_path: pytorch_geometric/torch_geometric/nn/summary.py
imports: from collections import defaultdict
  from typing import Any, List, Optional, Union
  import torch
  from torch.jit import ScriptModule
  from torch.nn import Module
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.nn.dense.linear import is_uninitialized_parameter
  from torch_geometric.typing import SparseTensor
label_desc: Use image node_id 1 for calling a global function with example usage: summary(model) and returns: make_table, hook