参考来源：
CSDN：GCN torchgeometric utils scatter方法源码解析及样例

源码及注释

源码位置：torch_geometric\utils\scatter.py
代码相当于在原有 torch_scatter 包的基础上进行二次开发，1.4.3 之后版本移除了，相当于只要理解 torch.scatter(…, reduce='add') 就好了。
下面是开发者说的原话：
“The scatter call got removed in one of the more recent versions, and you can now simply use torch.scatter(…, reduce=‘add’) for the same effect.”（scatter 调用在较新的版本之一中被删除，您现在可以简单地使用 torch.scatter(…, reduce=’add’) 来获得相同的效果。）

  1. import torch_scatter
  4. def scatter_(name, src, index, dim=0, dim_size=None):
  5.     r"""Aggregates all values from the :attr:`src` tensor at the indices
  6.     specified in the :attr:`index` tensor along the first dimension.
  7.     If multiple indices reference the same location, their contributions
  8.     are aggregated according to :attr:`name` (either :obj:`"add"`,
  9.     :obj:`"mean"` or :obj:`"max"`).
  11.     Args:
  12.         name (string): The aggregation to use (:obj:`"add"`, :obj:`"mean"`,
  13.             :obj:`"min"`, :obj:`"max"`).
  14.         src (Tensor): The source tensor.
  15.         index (LongTensor): The indices of elements to scatter.
  16.         dim (int, optional): The axis along which to index. (default: :obj:`0`)
  17.         dim_size (int, optional): Automatically create output tensor with size
  18.             :attr:`dim_size` in the first dimension. If set to :attr:`None`, a
  19.             minimal sized output tensor is returned. (default: :obj:`None`)
  21.     :rtype: :class:`Tensor`
  22.     """
  23.     # 例行断言
  24.     assert name in ['add', 'mean', 'min', 'max']
  26.     # 返回对象属性值
  27.     op = getattr(torch_scatter, 'scatter_{}'.format(name))
  28.     # 获取对象属性后返回值可直接使用，看具体案例
  29.     out = op(src, index, dim, None, dim_size)
  30.     out = out[0] if isinstance(out, tuple) else out
  32.     # 限制取最值后数据约束
  33.     if name == 'max':
  34.         out[out < -10000] = 0
  35.     elif name == 'min':
  36.         out[out > 10000] = 0
  38.     return out

torch_scatter 方法

结合官网和以下网址大家可以看得懂（看不懂没关系，多看几遍，我也是看了好几遍才懂的，不得不佩服其设计），这里不做过多解释。
语雀：torch.tensor.scatter() 和 torch.tensor.scatter_() 函数
Pytorch 文档：https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_.html

scatter_ 方法解析

常规运行

在理解上述原有开发包的基础上，我们再来看一下此方法的作用。按传入方式聚合节点信息。其是用于创建消息传递层的基类。
为了便于理解，这里给出一个具体例子：
假设我们有一个图，图中有四个节点，节点特征维度为 3。
其关系如下图所示（为方便理解我画成了单向图）：
image.png

  1. import torch
  2. import torch_scatter
  5. def scatter_(name, src, index, dim=0, dim_size=None):
  6.     r"""Aggregates all values from the :attr:`src` tensor at the indices
  7.     specified in the :attr:`index` tensor along the first dimension.
  8.     If multiple indices reference the same location, their contributions
  9.     are aggregated according to :attr:`name` (either :obj:`"add"`,
  10.     :obj:`"mean"` or :obj:`"max"`).
  12.     Args:
  13.         name (string): The aggregation to use (:obj:`"add"`, :obj:`"mean"`,
  14.             :obj:`"min"`, :obj:`"max"`).
  15.         src (Tensor): The source tensor.
  16.         index (LongTensor): The indices of elements to scatter.
  17.         dim (int, optional): The axis along which to index. (default: :obj:`0`)
  18.         dim_size (int, optional): Automatically create output tensor with size
  19.             :attr:`dim_size` in the first dimension. If set to :attr:`None`, a
  20.             minimal sized output tensor is returned. (default: :obj:`None`)
  22.     :rtype: :class:`Tensor`
  23.     """
  24.     # 例行断言
  25.     assert name in ['add', 'mean', 'min', 'max']
  27.     # 返回对象属性值
  28.     op = getattr(torch_scatter, 'scatter_{}'.format(name))
  29.     # 获取对象属性后返回值可直接使用，看具体案例
  30.     out = op(src, index, dim, None, dim_size)
  31.     # out = src.clone()
  32.     # out = op(src, index, dim, out, dim_size)
  33.     out = out[0] if isinstance(out, tuple) else out
  35.     # 限制取最值后数据约束
  36.     if name == 'max':
  37.         out[out < -10000] = 0
  38.     elif name == 'min':
  39.         out[out > 10000] = 0
  41.     return out
  44. # src 节点特征
  45. a = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [3.0, 3.0, 3.0], [4.0, 4.0, 4.0]])
  46. # index 索引
  47. b = torch.tensor([2, 3, 0, 1])
  48. # dim 维度信息
  49. c = 1
  50. # dim_size 返回的维度
  51. e = 4
  52. print(scatter_('add', a, b, dim=0, dim_size=e))

结果：

  1. tensor([[3., 3., 3.],
  2.         [4., 4., 4.],
  3.         [1., 1., 1.],
  4.         [2., 2., 2.]])

可以看到他进行了一次消息传递，以节点 1 为例，他接收了来自节点 3 的消息，所以特征变为 [4., 4., 4.]。

节点特征维度

此样例说明节点信息聚合时，各个维度相互独立。

  1. a = torch.tensor([[1.1, 1.2, 1.3], [2.1, 2.2, 2.3], [3.1, 3.2, 3.3], [4.1, 4.2, 4.3]])
  2. b = torch.tensor([2, 3, 0, 1])
  3. c = 1
  4. e = 4
  5. print(scatter_('add', a, b, dim=0, dim_size=e))
  8. """output:
  9. tensor([[3.1000, 3.2000, 3.3000],
  10.         [4.1000, 4.2000, 4.3000],
  11.         [1.1000, 1.2000, 1.3000],
  12.         [2.1000, 2.2000, 2.3000]])
  13. """

dim_size（运算后返回的维度）参数

只修改 dim_size，可以看到返回维度增加。

  1. a = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [3.0, 3.0, 3.0], [4.0, 4.0, 4.0]])
  2. b = torch.tensor([2, 3, 0, 1])
  3. c = 1
  4. e = 8
  5. print(scatter_('add', a, b, dim=0, dim_size=e))
  8. """output:
  9. tensor([[3., 3., 3.],
  10.         [4., 4., 4.],
  11.         [1., 1., 1.],
  12.         [2., 2., 2.],
  13.         [0., 0., 0.],
  14.         [0., 0., 0.],
  15.         [0., 0., 0.],
  16.         [0., 0., 0.]])
  17. """

index（索引）参数

这里相当于修改了节点的指向，修改后图：
image.png

  1. a = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [3.0, 3.0, 3.0], [4.0, 4.0, 4.0]])
  2. b = torch.tensor([2, 0, 0, 1])
  3. c = 1
  4. e = 4
  5. print(scatter_('add', a, b, dim=0, dim_size=e))
  8. """output:
  9. tensor([[5., 5., 5.],
  10.         [4., 4., 4.],
  11.         [1., 1., 1.],
  12.         [0., 0., 0.]])
  13. """

因为这里使用的 'add' 方法，所以同时传递到一个节点的信息被累加。

name（聚合方式）

很显然，当我把聚合方式改为均值，除了节点1有两个值传入被平均，其他保持不变；同理最值一样道理。

  1. a = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [3.0, 3.0, 3.0], [4.0, 4.0, 4.0]])
  2. b = torch.tensor([2, 0, 0, 1])
  3. c = 1
  4. e = 4
  5. print(scatter_('mean', a, b, dim=0, dim_size=e))
  8. """output:
  9. tensor([[2.5000, 2.5000, 2.5000],
  10.         [4.0000, 4.0000, 4.0000],
  11.         [1.0000, 1.0000, 1.0000],
  12.         [0.0000, 0.0000, 0.0000]])
  13. """

1.4.3 之后版本

源码位置：Lib\site-packages\torch_scatter\scatter.py

  1. from typing import Optional, Tuple
  2. import torch
  3. from .utils import broadcast
  6. def scatter_sum(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  7.                 out: Optional[torch.Tensor] = None,
  8.                 dim_size: Optional[int] = None) -> torch.Tensor:
  9.     index = broadcast(index, src, dim)
  10.     if out is None:
  11.         size = list(src.size())
  12.         if dim_size is not None:
  13.             size[dim] = dim_size
  14.         elif index.numel() == 0:
  15.             size[dim] = 0
  16.         else:
  17.             size[dim] = int(index.max()) + 1
  18.         out = torch.zeros(size, dtype=src.dtype, device=src.device)
  19.         return out.scatter_add_(dim, index, src)
  20.     else:
  21.         return out.scatter_add_(dim, index, src)
  24. def scatter_add(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  25.                 out: Optional[torch.Tensor] = None,
  26.                 dim_size: Optional[int] = None) -> torch.Tensor:
  27.     return scatter_sum(src, index, dim, out, dim_size)
  30. def scatter_mul(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  31.                 out: Optional[torch.Tensor] = None,
  32.                 dim_size: Optional[int] = None) -> torch.Tensor:
  33.     return torch.ops.torch_scatter.scatter_mul(src, index, dim, out, dim_size)
  36. def scatter_mean(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  37.                  out: Optional[torch.Tensor] = None,
  38.                  dim_size: Optional[int] = None) -> torch.Tensor:
  40.     out = scatter_sum(src, index, dim, out, dim_size)
  41.     dim_size = out.size(dim)
  43.     index_dim = dim
  44.     if index_dim < 0:
  45.         index_dim = index_dim + src.dim()
  46.     if index.dim() <= index_dim:
  47.         index_dim = index.dim() - 1
  49.     ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
  50.     count = scatter_sum(ones, index, index_dim, None, dim_size)
  51.     count[count < 1] = 1
  52.     count = broadcast(count, out, dim)
  53.     if out.is_floating_point():
  54.         out.true_divide_(count)
  55.     else:
  56.         out.floor_divide_(count)
  57.     return out
  60. def scatter_min(
  61.         src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  62.         out: Optional[torch.Tensor] = None,
  63.         dim_size: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
  64.     return torch.ops.torch_scatter.scatter_min(src, index, dim, out, dim_size)
  67. def scatter_max(
  68.         src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  69.         out: Optional[torch.Tensor] = None,
  70.         dim_size: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
  71.     return torch.ops.torch_scatter.scatter_max(src, index, dim, out, dim_size)
  74. def scatter(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
  75.             out: Optional[torch.Tensor] = None, dim_size: Optional[int] = None,
  76.             reduce: str = "sum") -> torch.Tensor:
  77.     r"""
  78.     |
  80.     .. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
  81.             master/docs/source/_figures/add.svg?sanitize=true
  82.         :align: center
  83.         :width: 400px
  85.     |
  87.     Reduces all values from the :attr:`src` tensor into :attr:`out` at the
  88.     indices specified in the :attr:`index` tensor along a given axis
  89.     :attr:`dim`.
  90.     For each value in :attr:`src`, its output index is specified by its index
  91.     in :attr:`src` for dimensions outside of :attr:`dim` and by the
  92.     corresponding value in :attr:`index` for dimension :attr:`dim`.
  93.     The applied reduction is defined via the :attr:`reduce` argument.
  95.     Formally, if :attr:`src` and :attr:`index` are :math:`n`-dimensional
  96.     tensors with size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})`
  97.     and :attr:`dim` = `i`, then :attr:`out` must be an :math:`n`-dimensional
  98.     tensor with size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`.
  99.     Moreover, the values of :attr:`index` must be between :math:`0` and
  100.     :math:`y - 1`, although no specific ordering of indices is required.
  101.     The :attr:`index` tensor supports broadcasting in case its dimensions do
  102.     not match with :attr:`src`.
  104.     For one-dimensional tensors with :obj:`reduce="sum"`, the operation
  105.     computes
  107.     .. math::
  108.         \mathrm{out}_i = \mathrm{out}_i + \sum_j~\mathrm{src}_j
  110.     where :math:`\sum_j` is over :math:`j` such that
  111.     :math:`\mathrm{index}_j = i`.
  113.     .. note::
  115.         This operation is implemented via atomic operations on the GPU and is
  116.         therefore **non-deterministic** since the order of parallel operations
  117.         to the same value is undetermined.
  118.         For floating-point variables, this results in a source of variance in
  119.         the result.
  121.     :param src: The source tensor.
  122.     :param index: The indices of elements to scatter.
  123.     :param dim: The axis along which to index. (default: :obj:`-1`)
  124.     :param out: The destination tensor.
  125.     :param dim_size: If :attr:`out` is not given, automatically create output
  126.         with size :attr:`dim_size` at dimension :attr:`dim`.
  127.         If :attr:`dim_size` is not given, a minimal sized output tensor
  128.         according to :obj:`index.max() + 1` is returned.
  129.     :param reduce: The reduce operation (:obj:`"sum"`, :obj:`"mul"`,
  130.         :obj:`"mean"`, :obj:`"min"` or :obj:`"max"`). (default: :obj:`"sum"`)
  132.     :rtype: :class:`Tensor`
  134.     .. code-block:: python
  136.         from torch_scatter import scatter
  138.         src = torch.randn(10, 6, 64)
  139.         index = torch.tensor([0, 1, 0, 1, 2, 1])
  141.         # Broadcasting in the first and last dim.
  142.         out = scatter(src, index, dim=1, reduce="sum")
  144.         print(out.size())
  146.     .. code-block::
  148.         torch.Size([10, 3, 64])
  149.     """
  150.     if reduce == 'sum' or reduce == 'add':
  151.         return scatter_sum(src, index, dim, out, dim_size)
  152.     if reduce == 'mul':
  153.         return scatter_mul(src, index, dim, out, dim_size)
  154.     elif reduce == 'mean':
  155.         return scatter_mean(src, index, dim, out, dim_size)
  156.     elif reduce == 'min':
  157.         return scatter_min(src, index, dim, out, dim_size)[0]
  158.     elif reduce == 'max':
  159.         return scatter_max(src, index, dim, out, dim_size)[0]
  160.     else:
  161.         raise ValueError