pyg-nightly 2.7.0.dev20250603__tar.gz → 2.7.0.dev20250605__tar.gz

{pyg_nightly-2.7.0.dev20250603 → pyg_nightly-2.7.0.dev20250605}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250603
+Version: 2.7.0.dev20250605
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.7.0.dev20250603 → pyg_nightly-2.7.0.dev20250605}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250603"
+version="2.7.0.dev20250605"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]

{pyg_nightly-2.7.0.dev20250603 → pyg_nightly-2.7.0.dev20250605}/torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250603'
+__version__ = '2.7.0.dev20250605'
 
 __all__ = [
     'Index',

{pyg_nightly-2.7.0.dev20250603 → pyg_nightly-2.7.0.dev20250605}/torch_geometric/datasets/city.py

@@ -16,11 +16,12 @@ class CityNetwork(InMemoryDataset):
     a Large Graph Dataset and a Measurement"
     <https://arxiv.org/abs/2503.09008>`_ paper.
     The dataset contains four city networks: `paris`, `shanghai`, `la`,
-    and 'london', where nodes represent junctions and edges represent
-    directed road segments. The task is to predict each node's eccentricity
-    score, which is approximated based on its 16-hop neighborhood. The score
-    indicates how accessible one node is in the network, and is mapped to
-    10 quantiles for transductive classification. See the original
+    and `london`, where nodes represent junctions and edges represent
+    undirected road segments. The task is to predict each node's eccentricity
+    score, which is approximated based on its 16-hop neighborhood and naturally
+    requires long-range information. The score indicates how accessible one
+    node is in the network, and is mapped to 10 quantiles for transductive
+    classification. See the original
     `source code <https://github.com/LeonResearch/City-Networks>`_ for more
     details on the individual networks.
 

{pyg_nightly-2.7.0.dev20250603 → pyg_nightly-2.7.0.dev20250605}/torch_geometric/nn/conv/__init__.py

@@ -61,6 +61,7 @@ from .gps_conv import GPSConv
 from .antisymmetric_conv import AntiSymmetricConv
 from .dir_gnn_conv import DirGNNConv
 from .mixhop_conv import MixHopConv
+from .meshcnn_conv import MeshCNNConv
 
 import torch_geometric.nn.conv.utils  # noqa
 
@@ -131,6 +132,7 @@ __all__ = [
     'AntiSymmetricConv',
     'DirGNNConv',
     'MixHopConv',
+    'MeshCNNConv',
 ]
 
 classes = __all__

pyg_nightly-2.7.0.dev20250605/torch_geometric/nn/conv/meshcnn_conv.py

@@ -0,0 +1,491 @@
+# The below is to suppress the warning on torch.nn.conv.MeshCNNConv::update
+# pyright: reportIncompatibleMethodOverride=false
+from typing import Optional
+from warnings import warn
+
+import torch
+from torch.nn import Linear, Module, ModuleList
+
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.typing import Tensor
+
+
+class MeshCNNConv(MessagePassing):
+    r"""The convolutional layer introduced by the paper
+    `"MeshCNN: A Network With An Edge" <https://arxiv.org/abs/1809.05910>`_.
+
+    Recall that, given a set of categories :math:`C`,
+    MeshCNN is a function that takes as its input
+    a triangular mesh
+    :math:`\mathcal{m} = (V, F) \in \mathbb{R}^{|V| \times 3} \times
+    \{0,...,|V|-1\}^{3 \times |F|}`, and returns as its output
+    a :math:`|C|`-dimensional vector, whose :math:`i` th component denotes
+    the probability of the input mesh belonging to category :math:`c_i \in C`.
+
+    Let :math:`X^{(k)} \in \mathbb{R}^{|E| \times \text{Dim-Out}(k)}`
+    denote the output value of the prior (e.g. :math:`k` th )
+    layer of our neural network. The :math:`i` th row of :math:`X^{(k)}` is a
+    :math:`\text{Dim-Out}(k)`-dimensional vector that represents the features
+    computed by the :math:`k` th layer for edge :math:`e_i` of the input mesh
+    :math:`\mathcal{m}`. Let :math:`A \in \{0, ..., |E|-1\}^{2 \times 4*|E|}`
+    denote the *edge adjacency* matrix of our input mesh :math:`\mathcal{m}`.
+    The :math:`j` th column of :math:`A` returns a pair of indices
+    :math:`k,l \in \{0,...,|E|-1\}`, which means that edge
+    :math:`e_k` is adjacent to edge :math:`e_l`
+    in our input mesh :math:`\mathcal{m}`.
+    The definition of edge adjacency in a triangular
+    mesh is illustrated in Figure 1.
+    In a triangular
+    mesh, each edge :math:`e_i` is expected to be adjacent to exactly :math:`4`
+    neighboring edges, hence the number of columns of :math:`A`: :math:`4*|E|`.
+    We write *the neighborhood* of edge :math:`e_i` as
+    :math:`\mathcal{N}(i) = (a(i), b(i), c(i), d(i))` where
+
+    1. :math:`a(i)` denotes the index of the *first* counter-clockwise
+    edge of the face *above* :math:`e_i`.
+
+    2. :math:`b(i)` denotes the index of the *second* counter-clockwise
+    edge of the face *above* :math:`e_i`.
+
+    3. :math:`c(i)` denotes the index of the *first* counter-clockwise edge
+    of the face *below* :math:`e_i`.
+
+    4. :math:`d(i)` denotes the index of the *second*
+    counter-clockwise edge of the face *below* :math:`e_i`.
+
+    .. figure:: ../_figures/meshcnn_edge_adjacency.svg
+        :align: center
+        :width: 80%
+
+        **Figure 1:** The neighbors of edge :math:`\mathbf{e_1}`
+        are :math:`\mathbf{e_2}, \mathbf{e_3}, \mathbf{e_4}` and
+        :math:`\mathbf{e_5}`, respectively.
+        We write this as
+        :math:`\mathcal{N}(1) = (a(1), b(1), c(1), d(1)) = (2, 3, 4, 5)`
+
+
+    Because of this ordering constrait, :obj:`MeshCNNConv` **requires
+    that the columns of** :math:`A`
+    **be ordered in the following way**:
+
+    .. math::
+        &A[:,0] = (0, \text{The index of the "a" edge for edge } 0) \\
+        &A[:,1] = (0, \text{The index of the "b" edge for edge } 0) \\
+        &A[:,2] = (0, \text{The index of the "c" edge for edge } 0) \\
+        &A[:,3] = (0, \text{The index of the "d" edge for edge } 0) \\
+        \vdots \\
+        &A[:,4*|E|-4] =
+            \bigl(|E|-1,
+                a\bigl(|E|-1\bigr)\bigr) \\
+        &A[:,4*|E|-3] =
+            \bigl(|E|-1,
+                b\bigl(|E|-1\bigr)\bigr) \\
+        &A[:,4*|E|-2] =
+            \bigl(|E|-1,
+                c\bigl(|E|-1\bigr)\bigr) \\
+        &A[:,4*|E|-1] =
+            \bigl(|E|-1,
+                d\bigl(|E|-1\bigr)\bigr)
+
+
+    Stated a bit more compactly, for every edge :math:`e_i` in the input mesh,
+    :math:`A`, should have the following entries
+
+    .. math::
+        A[:, 4*i] &= (i, a(i)) \\
+        A[:, 4*i + 1] &= (i, b(i)) \\
+        A[:, 4*i + 2] &= (i, c(i)) \\
+        A[:, 4*i + 3] &= (i, d(i))
+
+    To summarize so far, we have defined 3 things:
+
+    1. The activation of the prior (e.g. :math:`k` th) layer,
+    :math:`X^{(k)} \in \mathbb{R}^{|E| \times \text{Dim-Out}(k)}`
+
+    2. The edge adjacency matrix and the definition of edge adjacency.
+    :math:`A \in \{0,...,|E|-1\}^{2 \times 4*|E|}`
+
+    3. The ways the columns of :math:`A` must be ordered.
+
+
+
+    We are now finally able to define the :obj:`MeshCNNConv` class/layer.
+    In the following definition
+    we assume :obj:`MeshCNNConv` is at the :math:`k+1` th layer of our
+    neural network.
+
+    The :obj:`MeshCNNConv` layer is a function,
+
+    .. math::
+        \text{MeshCNNConv}^{(k+1)}(X^{(k)}, A) = X^{(k+1)},
+
+    that, given the prior layer's output
+    :math:`X^{(k)} \in \mathbb{R}^{|E| \times \text{Dim-Out}(k)}`
+    and the edge adjacency matrix :math:`A`
+    of the input mesh (graph) :math:`\mathcal{m}` ,
+    returns a new edge feature tensor
+    :math:`X^{(k+1)} \in \mathbb{R}^{|E| \times \text{Dim-Out}(k+1)}`,
+    where the :math:`i` th row of :math:`X^{(k+1)}`, denoted by
+    :math:`x^{(k+1)}_i`,
+    represents the :math:`\text{Dim-Out}(k+1)`-dimensional feature vector
+    of edge :math:`e_i`, **and is defined as follows**:
+
+    .. math::
+        x^{(k+1)}_i &= W^{(k+1)}_0 x^{(k)}_i \\
+        &+ W^{(k+1)}_1 \bigl| x^{(k)}_{a(i)} - x^{(k)}_{c(i)} \bigr| \\
+        &+ W^{(k+1)}_2 \bigl( x^{(k)}_{a(i)} + x^{(k)}_{c(i)} \bigr) \\
+        &+ W^{(k+1)}_3 \bigl| x^{(k)}_{b(i)} - x^{(k)}_{d(i)} \bigr| \\
+        &+ W^{(k+1)}_4 \bigl( x^{(k)}_{b(i)} + x^{(k)}_{d(i)} \bigr).
+
+    :math:`W_0^{(k+1)},W_1^{(k+1)},W_2^{(k+1)},W_3^{(k+1)}, W_4^{(k+1)}
+    \in \mathbb{R}^{\text{Dim-Out}(k+1) \times \text{Dim-Out}(k)}`
+    are trainable linear functions (i.e. "the weights" of this layer).
+    :math:`x_i` is the :math:`\text{Dim-Out}(k)`-dimensional feature of
+    edge :math:`e_i` vector computed by the prior (e.g. :math:`k`) th layer.
+    :math:`x^{(k)}_{a(i)}, x^{(k)}_{b(i)}, x^{(k)}_{c(i)}`, and
+    :math:`x^{(k)}_{d(i)}` are the :math:`\text{Dim-Out}(k)`-feature vectors,
+    computed in the :math:`k` th layer, that are associated with the :math:`4`
+    neighboring edges of :math:`e_i`.
+
+
+    Args:
+        in_channels (int): Corresonds to :math:`\text{Dim-Out}(k)`
+            in the above overview. This
+            represents the output dimension of the prior layer. For the given
+            input mesh :math:`\mathcal{m} = (V, F)`, the prior layer is
+            expected to output a
+            :math:`X \in \mathbb{R}^{|E| \times \textit{in_channels}}`
+            feature matrix.
+            Assuming the instance of this class
+            is situated at layer :math:`k+1`, we write that
+            :math:`X^{(k)} \in \mathbb{R}^{|E| \times \textit{in_channels}}`.
+        out_channels (int): Corresponds to :math:`\text{Dim-Out}(k+1)` in the
+            above overview. This represents the output dimension of this layer.
+            Assuming the instance of this class
+            is situated at layer :math:`k+1`, we write that
+            :math:`X^{(k+1)}
+            \in \mathbb{R}^{|E| \times \textit{out_channels}}`.
+        kernels (torch.nn.ModuleList, optional): A list of length of 5,
+            where each
+            element is a :class:`torch.nn.module` (i.e a neural network),
+            that each MUST take as input a vector
+            of dimension :`obj:in_channels` and return a vector of dimension
+            :obj:`out_channels`. In particular,
+            `obj:kernels[0]` is :math:`W^{(k+1)}_0` in the above overview
+            (see :obj:`MeshCNNConv`), `obj:kernels[1]` is :math:`W^{(k+1)}_1`,
+            `obj:kernels[2]` is :math:`W^{(k+1)}_2`,
+            `obj:kernels[3]` is :math:`W^{(k+1)}_3`
+            `obj:kernels[4]` is :math:`W^{(k+1)}_4`.
+            Note that this input is optional, in which case
+            each of the 5 elements in the kernels will be a linear
+            neural network :class:`torch.nn.modules.Linear`
+            correctly configured to take as input
+            :attr:`in_channels`-dimensional vectors and return
+            a vector of dimensions :attr:`out_channels`.
+
+    Discussion:
+        The key difference that seperates :obj:`MeshCNNConv` from a traditional
+        message passing graph neural network is that :obj:`MeshCNNConv`
+        requires the set of neighbors for a node
+        :math:`\mathcal{N}(u) = (v_1, v_2, ...)`
+        to *be an ordered set* (i.e. a tuple). In
+        fact, :obj:`MeshCNNConv` goes further, requiring
+        that :math:`\mathcal{N}(u)` always return a set of size :math:`4`.
+        This is different to most message passing graph neural networks,
+        which assume that :math:`\mathcal{N}(u) = \{v_1, v_2, ...\}` returns an
+        ordered set. This lends :obj:`MeshCNNConv` more expressive power,
+        at the cost of no longer being permutation invariant to
+        :math:`\mathbb{S}_4`. Put more plainly, in tradition message passing
+        GNNs, the network is *unable* to distinguish one neighboring node
+        from another.
+        In constrast, in :obj:`MeshCNNConv`, each of the 4 neighbors has a
+        "role", either the "a", "b", "c", or "d" neighbor. We encode this fact
+        by requiring that :math:`\mathcal{N}` return the 4-tuple,
+        where the first component is the "a" neighbor, and so on.
+
+        To summarize this comparison, it may re-define
+        :obj:`MeshCNNConv` in terms of :math:`\text{UPDATE}` and
+        :math:`\text{AGGREGATE}`
+        functions, which is a general way to define a traditional GNN layer.
+        If we let :math:`x_i^{(k+1)}`
+        denote the output of a GNN layer for node :math:`i` at
+        layer :math:`k+1`, and let
+        :math:`\mathcal{N}(i)` denote the set of nodes adjacent
+        to node :math:`i`,
+        then we can describe the :math:`k+1` th layer as traditional GNN
+        as
+
+        .. math::
+            x_i^{(k+1)} = \text{UPDATE}^{(k+1)}\bigl(x^{(k)}_i,
+            \text{AGGREGATE}^{(k+1)}\bigl(\mathcal{N}(i)\bigr)\bigr).
+
+        Here, :math:`\text{UPDATE}^{(k+1)}` is a function of :math:`2`
+        :math:`\text{Dim-Out}(k)`-dimensional vectors, and returns a
+        :math:`\text{Dim-Out}(k+1)`-dimensional vector.
+        :math:`\text{AGGREGATE}^{(k+1)}` function
+        is a function of a *unordered set*
+        of nodes that are neighbors of node :math:`i`, as defined by
+        :math:`\mathcal{N}(i)`. Usually the size of this set varies across
+        different nodes :math:`i`, and one of the most basic examples
+        of such a function is the "sum aggregation", defined as
+        :math:`\text{AGGREGATE}^{(k+1)}(\mathcal{N}(i)) =
+        \sum_{j \in \mathcal{N}(i)} x^{(k)}_j`.
+        See
+        :class:`SumAggregation <torch_geometric.nn.aggr.basic.SumAggregation>`
+        for more.
+
+        In contrast, while :obj:`MeshCNNConv` 's :math:`\text{UPDATE}`
+        function follows
+        a tradition GNN, its :math:`\text{AGGREGATE}` is a function of a tuple
+        (i.e. an ordered set) of neighbors
+        rather than a unordered set of neighbors.
+        In particular, while the :math:`\text{UPDATE}`
+        function of :obj:`MeshCNNConv` for :math:`e_i` is
+
+        .. math::
+            x_i^{(k+1)} = \text{UPDATE}^{(k+1)}(x_i^{(k)}, s_i^{(k+1)})
+            = W_0^{(k+1)}x_i^{(k)} + s_i^{(k+1)},
+
+        in contrast, :obj:`MeshCNNConv` 's :math:`\text{AGGREGATE}` function is
+
+        .. math::
+            s_i^{(k+1)} = \text{AGGREGATE}^{(k+1)}(A, B, C, D)
+            &= W_1^{(k+1)}\bigl|A - C \bigr| \\
+            &= W_2^{(k+1)}\bigl(A + C \bigr) \\
+            &= W_3^{(k+1)}\bigl|B - D \bigr| \\
+            &= W_4^{(k+1)}\bigl(B + D \bigr),
+
+        where :math:`A=x_{a(i)}^{(k)}, B=x_{b(i)}^{(k)}, C=x_{c(i)}^{(k)},`
+        and :math:`D=x_{d(i)}^{(k)}`.
+
+        ..
+
+            The :math:`i` th row of
+            :math:`V \in \mathbb{R}^{|V| \times 3}`
+            holds the cartesian :math:`xyz`
+            coordinates for node :math:`v_i` in the mesh, and the :math:`j` th
+            column in :math:`F \in \{1,...,|V|\}^{3 \times |V|}`
+            holds the :math:`3` indices
+            :math:`(k,l,m)` that correspond to the :math:`3` nodes
+            :math:`(v_k, v_l, v_m)` that construct face :math:`j` of the mesh.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernels: Optional[ModuleList] = None):
+        super().__init__(aggr='add')
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if kernels is None:
+            self.kernels = ModuleList(
+                [Linear(in_channels, out_channels) for _ in range(5)])
+
+        else:
+            # ensures kernels is properly formed, otherwise throws
+            # the appropriate error.
+            self._assert_kernels(kernels)
+            self.kernels = kernels
+
+    def forward(self, x: Tensor, edge_index: Tensor):
+        r"""Forward pass.
+
+        Args:
+            x(torch.Tensor): :math:`X^{(k)} \in
+                \mathbb{R}^{|E| \times \textit{in_channels}}`.
+                The edge feature tensor returned by the prior layer
+                (e.g. :math:`k`). The tensor is of shape
+                :math:`|E| \times \text{Dim-Out}(k)`, or equivalently,
+                :obj:`(|E|, self.in_channels)`.
+
+            edge_index(torch.Tensor):
+                :math:`A \in \{0,...,|E|-1\}^{2 \times 4*|E|}`.
+                The edge adjacency tensor of the networks input mesh
+                :math:`\mathcal{m} = (V, F)`. The edge adjacency tensor
+                **MUST** have the following form:
+
+                .. math::
+                    &A[:,0] = (0,
+                        \text{The index of the "a" edge for edge } 0) \\
+                    &A[:,1] = (0,
+                        \text{The index of the "b" edge for edge } 0) \\
+                    &A[:,2] = (0,
+                        \text{The index of the "c" edge for edge } 0) \\
+                    &A[:,3] = (0,
+                        \text{The index of the "d" edge for edge } 0) \\
+                    \vdots \\
+                    &A[:,4*|E|-4] =
+                        \bigl(|E|-1,
+                            a\bigl(|E|-1\bigr)\bigr) \\
+                    &A[:,4*|E|-3] =
+                        \bigl(|E|-1,
+                            b\bigl(|E|-1\bigr)\bigr) \\
+                    &A[:,4*|E|-2] =
+                        \bigl(|E|-1,
+                            c\bigl(|E|-1\bigr)\bigr) \\
+                    &A[:,4*|E|-1] =
+                        \bigl(|E|-1,
+                            d\bigl(|E|-1\bigr)\bigr)
+
+                See :obj:`MeshCNNConv` for what
+                "index of the 'a'(b,c,d) edge for edge i" means, and also
+                for the general definition of edge adjacency in MeshCNN.
+                These definitions are also provided in the
+                `paper <https://arxiv.org/abs/1809.05910>`_ itself.
+
+        Returns:
+           torch.Tensor:
+           :math:`X^{(k+1)} \in \mathbb{R}^{|E| \times \textit{out_channels}}`.
+           The edge feature tensor for this (e.g. the :math:`k+1` th) layer.
+           The :math:`i` th row of :math:`X^{(k+1)}` is computed according
+           to the formula
+
+            .. math::
+                x^{(k+1)}_i &= W^{(k+1)}_0 x^{(k)}_i \\
+                &+ W^{(k+1)}_1 \bigl| x^{(k)}_{a(i)} - x^{(k)}_{c(i)} \bigr| \\
+                &+ W^{(k+1)}_2 \bigl( x^{(k)}_{a(i)} + x^{(k)}_{c(i)} \bigr) \\
+                &+ W^{(k+1)}_3 \bigl| x^{(k)}_{b(i)} - x^{(k)}_{d(i)} \bigr| \\
+                &+ W^{(k+1)}_4 \bigl( x^{(k)}_{b(i)} + x^{(k)}_{d(i)} \bigr),
+
+            where :math:`W_0^{(k+1)},W_1^{(k+1)},
+            W_2^{(k+1)},W_3^{(k+1)}, W_4^{(k+1)}
+            \in \mathbb{R}^{\text{Dim-Out}(k+1) \times \text{Dim-Out}(k)}`
+            are the trainable linear functions (i.e. the trainable
+            "weights") of this layer, and
+            :math:`x^{(k)}_{a(i)}, x^{(k)}_{b(i)}, x^{(k)}_{c(i)}`,
+            :math:`x^{(k)}_{d(i)}` are the
+            :math:`\text{Dim-Out}(k)`-dimensional edge feature vectors
+            computed by the prior (:math:`k` th) layer,
+            that are associated with the :math:`4`
+            neighboring edges of :math:`e_i`.
+
+        """
+        return self.propagate(edge_index, x=x)
+
+    def message(self, x_j: Tensor) -> Tensor:
+        r"""The messaging passing step of :obj:`MeshCNNConv`.
+
+
+        Args:
+          x_j: A :obj:`[4*|E|, num_node_features]` tensor.
+          Its ith row holds the value
+            stored by the source node in the previous layer of edge i.
+
+        Returns:
+            A :obj:`[|E|, num_node_features]` tensor,
+            whose ith row will be the value
+            that the target node of edge i will receive.
+        """
+        # The following variables names are taken from the paper
+        # MeshCNN computes the features associated with edge
+        # e by (|a - c|, a + c, |b - c|, b + c), where a, b, c, d are the
+        # neighboring edges of e, a being the 1 edge of the upper face,
+        # b being the second edge of the upper face, c being the first edge
+        # of the lower face,
+        # and d being the second edge of the lower face of the input Mesh
+
+        # TODO: It is unclear  if view is faster. If it is not,
+        # then we should prefer the strided method commented out below
+
+        E4, in_channels = x_j.size()  # E4 = 4|E|, i.e. num edges in line graph
+        # Option 1
+        n_a = x_j[0::4]  # shape: |E| x in_channels
+        n_b = x_j[1::4]  # shape: |E| x in_channels
+        n_c = x_j[2::4]  # shape: |E| x in_channels
+        n_d = x_j[3::4]  # shape: |E| x in_channels
+        m = torch.empty(E4, self.out_channels)
+        m[0::4] = self.kernels[1].forward(torch.abs(n_a - n_c))
+        m[1::4] = self.kernels[2].forward(n_a + n_c)
+        m[2::4] = self.kernels[3].forward(torch.abs(n_b - n_d))
+        m[3::4] = self.kernels[4].forward(n_b + n_d)
+        return m
+
+        # Option 2
+        # E4, in_channels = x_j.size()
+        # E = E4 // 4
+        # x_j = x_j.view(E, 4, in_channels)  # shape: (|E| x 4 x in_channels)
+        # n_a, n_b, n_c, n_d = x_j.unbind(
+        #     dim=1)  # shape: (4 x |E| x in_channels)
+        # m = torch.stack(
+        #     [
+        #         (n_a - n_c).abs(),  # shape: |E| x in_channels
+        #         n_a + n_c,
+        #         (n_b - n_d).abs(),
+        #         n_b + n_d,
+        #     ],
+        #     dim=1)  # shape: (|E| x 4 x in_channels)
+        # m.view(E4, in_channels)  # shape 4*|E| x in_channels
+        # return m
+
+    def update(self, inputs: Tensor, x: Tensor) -> Tensor:
+        r"""The UPDATE step, in reference to the UPDATE and AGGREGATE
+        formulation of message passing convolution.
+
+        Args:
+           inputs(torch.Tensor): The :attr:`in_channels`-dimensional vector
+            returned by aggregate.
+           x(torch.Tensor): :math:`X^{(k)}`. The original inputs to this layer.
+
+        Returns:
+            torch.Tensor: :math:`X^{(k+1)}`. The output of this layer, which
+            has shape :obj:`(|E|, out_channels)`.
+        """
+        return self.kernels[0].forward(x) + inputs
+
+    def _assert_kernels(self, kernels: ModuleList):
+        r"""Ensures that :obj:`kernels` is a list of 5 :obj:`torch.nn.Module`
+        modules (i.e. networks). In addition, it also ensures that each network
+        takes in input of dimension :attr:`in_channels`, and returns output
+        of dimension :attr:`out_channels`.
+        This method throws an error otherwise.
+
+        .. warn::
+            This method throws an error if :obj:`kernels` is
+            not valid. (Otherwise this method returns nothing)
+
+        """
+        assert isinstance(kernels, ModuleList), \
+            f"Parameter 'kernels' must be a \
+            torch.nn.module.ModuleList with 5 memebers, but we got \
+            {type(kernels)}."
+
+        assert len(kernels) == 5, "Parameter 'kernels' must be a \
+            torch.nn.module.ModuleList of with exactly 5 members"
+
+        for i, network in enumerate(kernels):
+            assert isinstance(network, Module), \
+                f"kernels[{i}] must be torch.nn.Module, got \
+                {type(network)}"
+            if not hasattr(network, "in_channels") and \
+                    not hasattr(network, "in_features"):
+                warn(f"kernel[{i}] does not have attribute \
+                            'in_channels' nor 'out_features'. The \
+                            network must take as input a \
+                            {self.in_channels}-dimensional tensor. \
+                            Still, assuming user configured \
+                            correctly. Continuing..")
+            else:
+                input_dimension = getattr(network, "in_channels",
+                                          network.in_features)
+                assert input_dimension == self.in_channels, f"The input \
+                dimension of the neural network in kernel[{i}] must \
+                be \
+                equal to 'in_channels', but input_dimension = \
+                {input_dimension}, and \
+                self.in_channels={self.in_channels}."
+
+            if not hasattr(network, "out_channels") and \
+                    not hasattr(network, "out_features"):
+                warn(f"kernel[{i}] does not have attribute \
+                                'in_channels' nor 'out_features'. The \
+                                network must take as input a \
+                                {self.in_channels}-dimensional tensor. \
+                                Still, assuming user configured \
+                                correctly. Continuing..")
+            else:
+                output_dimension = getattr(network, "out_channels",
+                                           network.out_features)
+                assert output_dimension == self.out_channels, f"The output \
+                    dimension of the neural network in kernel[{i}] must \
+                    be \
+                    equal to 'out_channels', but out_dimension = \
+                    {output_dimension}, and \
+                    self.out_channels={self.out_channels}."

{pyg_nightly-2.7.0.dev20250603 → pyg_nightly-2.7.0.dev20250605}/torch_geometric/utils/__init__.py

@@ -57,6 +57,7 @@ from .embedding import get_embeddings, get_embeddings_hetero
 from ._trim_to_layer import trim_to_layer
 from .ppr import get_ppr
 from ._train_test_split_edges import train_test_split_edges
+from .influence import total_influence
 
 __all__ = [
     'scatter',
@@ -149,6 +150,7 @@ __all__ = [
     'trim_to_layer',
     'get_ppr',
     'train_test_split_edges',
+    'total_influence',
 ]
 
 # `structured_negative_sampling_feasible` is a long name and thus destroys the