Repository: pyg-team/pytorch_geometric
Pull request #7924 (2023-08-24 06:01:13)

[Test] PyTorch 2.1 supports gradients on CSR matrices


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 968fa1e51..4946196c0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -87,6 +87,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Fixed the usage of `bias_initializer` in `HeteroLinear` ([#7923](https://github.com/pyg-team/pytorch_geometric/pull/7923))
 - Fixed broken links in `HGBDataset` ([#7907](https://github.com/pyg-team/pytorch_geometric/pull/7907))
 - Fixed an issue where `SetTransformerAggregation` produced `NaN` values for isolates nodes ([#7902](https://github.com/pyg-team/pytorch_geometric/pull/7902))
 - Fixed `model_summary` on modules with uninitialized parameters ([#7884](https://github.com/pyg-team/pytorch_geometric/pull/7884))
diff --git a/torch_geometric/data/collate.py b/torch_geometric/data/collate.py
index 134efb7a1..f217479b6 100644
--- a/torch_geometric/data/collate.py
+++ b/torch_geometric/data/collate.py
@@ -154,7 +154,7 @@ def _collate(
         if torch.utils.data.get_worker_info() is not None:
             # Write directly into shared memory to avoid an extra copy:
             numel = sum(value.numel() for value in values)
-            if torch_geometric.typing.WITH_PT2:
+            if torch_geometric.typing.WITH_PT20:
                 storage = elem.untyped_storage()._new_shared(
                     numel * elem.element_size(), device=elem.device)
             elif torch_geometric.typing.WITH_PT112:
diff --git a/torch_geometric/loader/utils.py b/torch_geometric/loader/utils.py
index 5da757319..a51fd7c3a 100644
--- a/torch_geometric/loader/utils.py
+++ b/torch_geometric/loader/utils.py
@@ -62,7 +62,7 @@ def index_select(
             size = list(value.shape)
             size[dim] = index.numel()
             numel = math.prod(size)
-            if torch_geometric.typing.WITH_PT2:
+            if torch_geometric.typing.WITH_PT20:
                 storage = value.untyped_storage()._new_shared(
                     numel * value.element_size())
             else:
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index 05900aac1..d7804c12d 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -249,8 +249,8 @@ class HeteroLinear(torch.nn.Module):
         r"""Resets all learnable parameters of the module."""
         reset_weight_(self.weight, self.in_channels,
                       self.kwargs.get('weight_initializer', None))
-        reset_weight_(self.bias, self.in_channels,
-                      self.kwargs.get('bias_initializer', None))
+        reset_bias_(self.bias, self.in_channels,
+                    self.kwargs.get('bias_initializer', None))
 
     def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
         r"""
diff --git a/torch_geometric/nn/models/rev_gnn.py b/torch_geometric/nn/models/rev_gnn.py
index 662457bed..2bf0ab15b 100644
--- a/torch_geometric/nn/models/rev_gnn.py
+++ b/torch_geometric/nn/models/rev_gnn.py
@@ -58,7 +58,7 @@ class InvertibleFunction(torch.autograd.Function):
         detached_outputs = tuple(element.detach_() for element in outputs)
 
         # Clear memory of node features:
-        if torch_geometric.typing.WITH_PT2:
+        if torch_geometric.typing.WITH_PT20:
             inputs[0].untyped_storage().resize_(0)
         else:  # pragma: no cover
             inputs[0].storage().resize_(0)
@@ -85,7 +85,7 @@ class InvertibleFunction(torch.autograd.Function):
             inputs_inverted = ctx.fn_inverse(*(outputs + inputs[1:]))
             if len(ctx.outputs) == 0:  # Clear memory from outputs:
                 for element in outputs:
-                    if torch_geometric.typing.WITH_PT2:
+                    if torch_geometric.typing.WITH_PT20:
                         element.untyped_storage().resize_(0)
                     else:  # pragma: no cover
                         element.storage().resize_(0)
@@ -94,7 +94,7 @@ class InvertibleFunction(torch.autograd.Function):
                 inputs_inverted = (inputs_inverted, )
 
             for elem_orig, elem_inv in zip(inputs, inputs_inverted):
-                if torch_geometric.typing.WITH_PT2:
+                if torch_geometric.typing.WITH_PT20:
                     elem_orig.untyped_storage().resize_(
                         int(np.prod(elem_orig.size())) *
                         elem_orig.element_size())
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 38a3a01fb..bdf3db12d 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -6,10 +6,11 @@ import numpy as np
 import torch
 from torch import Tensor
 
-WITH_PT2 = int(torch.__version__.split('.')[0]) >= 2
-WITH_PT111 = WITH_PT2 or int(torch.__version__.split('.')[1]) >= 11
-WITH_PT112 = WITH_PT2 or int(torch.__version__.split('.')[1]) >= 12
-WITH_PT113 = WITH_PT2 or int(torch.__version__.split('.')[1]) >= 13
+WITH_PT20 = int(torch.__version__.split('.')[0]) >= 2
+WITH_PT21 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 1
+WITH_PT111 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 11
+WITH_PT112 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 12
+WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
 
 if not hasattr(torch, 'sparse_csc'):
     torch.sparse_csc = -1
@@ -17,7 +18,7 @@ if not hasattr(torch, 'sparse_csc'):
 try:
     import pyg_lib  # noqa
     WITH_PYG_LIB = True
-    WITH_GMM = WITH_PT2 and hasattr(pyg_lib.ops, 'grouped_matmul')
+    WITH_GMM = WITH_PT20 and hasattr(pyg_lib.ops, 'grouped_matmul')
     WITH_SEGMM = hasattr(pyg_lib.ops, 'segment_matmul')
     if WITH_SEGMM and 'pytest' in sys.modules and torch.cuda.is_available():
         # NOTE `segment_matmul` is currently bugged on older NVIDIA cards which
diff --git a/torch_geometric/utils/spmm.py b/torch_geometric/utils/spmm.py
index b9dd77b99..d951b7e0e 100644
--- a/torch_geometric/utils/spmm.py
+++ b/torch_geometric/utils/spmm.py
@@ -43,7 +43,7 @@ def spmm(src: Adj, other: Tensor, reduce: str = "sum") -> Tensor:
         if src.nnz() == 0:
             return other.new_zeros(src.size(0), other.size(1))
 
-        if (torch_geometric.typing.WITH_PT2 and other.dim() == 2
+        if (torch_geometric.typing.WITH_PT20 and other.dim() == 2
                 and not src.is_cuda() and not src.requires_grad()):
             # Use optimized PyTorch `torch.sparse.mm` path:
             csr = src.to_torch_sparse_csr_tensor().to(other.dtype)
@@ -56,7 +56,7 @@ def spmm(src: Adj, other: Tensor, reduce: str = "sum") -> Tensor:
 
     # `torch.sparse.mm` only supports reductions on CPU for PyTorch>=2.0.
     # This will currently throw on error for CUDA tensors.
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
 
         if src.is_cuda and (reduce == 'min' or reduce == 'max'):
             raise NotImplementedError(f"`{reduce}` reduction is not yet "


Test patch:
diff --git a/test/nn/conv/test_gcn_conv.py b/test/nn/conv/test_gcn_conv.py
index 2efc52e27..63ba7da52 100644
--- a/test/nn/conv/test_gcn_conv.py
+++ b/test/nn/conv/test_gcn_conv.py
@@ -7,7 +7,7 @@ import torch_geometric.typing
 from torch_geometric.nn import GCNConv
 from torch_geometric.nn.conv.gcn_conv import gcn_norm
 from torch_geometric.testing import is_full_test
-from torch_geometric.typing import SparseTensor
+from torch_geometric.typing import WITH_PT21, SparseTensor
 from torch_geometric.utils import to_torch_coo_tensor, to_torch_csc_tensor
 
 
@@ -120,8 +120,8 @@ def test_gcn_norm_gradient(requires_grad, layout):
     if layout == torch.sparse_csr:
         adj = adj.to_sparse_csr()
 
-    # TODO Sparse CSR tensor does not yet inherit `requires_grad` from `value`.
-    if layout == torch.sparse_csr:
+    # TODO Sparse CSR tensor doesn't inherit `requires_grad` for PyTorch < 2.1.
+    if layout == torch.sparse_csr and not WITH_PT21:
         assert not gcn_norm(adj)[0].requires_grad
     else:
         assert adj.requires_grad == gcn_norm(adj)[0].requires_grad
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
index d193276dc..9c004b7b0 100644
--- a/test/nn/dense/test_linear.py
+++ b/test/nn/dense/test_linear.py
@@ -127,6 +127,17 @@ def test_hetero_linear(device):
     assert torch.allclose(jit(x, type_vec), out, atol=1e-3)
 
 
+def test_hetero_linear_initializer():
+    lin = HeteroLinear(
+        16,
+        32,
+        num_types=3,
+        weight_initializer='glorot',
+        bias_initializer='zeros',
+    )
+    assert torch.equal(lin.bias, torch.zeros_like(lin.bias))
+
+
 @withCUDA
 @pytest.mark.parametrize('use_segmm', [True, False])
 def test_hetero_linear_amp(device, use_segmm):
diff --git a/test/nn/models/test_rev_gnn.py b/test/nn/models/test_rev_gnn.py
index c114d97d3..d9d9a1756 100644
--- a/test/nn/models/test_rev_gnn.py
+++ b/test/nn/models/test_rev_gnn.py
@@ -22,7 +22,7 @@ def test_revgnn_forward_inverse(num_groups):
     h_o = h.clone().detach()
 
     out = conv(h, edge_index)
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
         assert h.untyped_storage().size() == 0
     else:
         assert h.storage().size() == 0
@@ -80,7 +80,7 @@ def test_revgnn_diable(num_groups):
     target.backward()
 
     # Memory will not be freed if disable:
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
         assert h.untyped_storage().size() == 4 * 4 * 32
     else:
         assert h.storage().size() == 4 * 32
diff --git a/test/utils/test_sparse.py b/test/utils/test_sparse.py
index 922e04cec..6c1af5e8d 100644
--- a/test/utils/test_sparse.py
+++ b/test/utils/test_sparse.py
@@ -134,7 +134,7 @@ def test_to_torch_csr_tensor():
     assert torch.allclose(coo.indices(), edge_index)
     assert torch.allclose(coo.values(), edge_weight)
 
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
         edge_attr = torch.randn(edge_index.size(1), 8)
         adj = to_torch_csr_tensor(edge_index, edge_attr)
         assert adj.size() == (4, 4, 8)
@@ -155,7 +155,7 @@ def test_to_torch_csc_tensor():
     assert adj.size() == (4, 4)
     assert adj.layout == torch.sparse_csc
     adj_coo = adj.to_sparse_coo().coalesce()
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
         assert torch.allclose(adj_coo.indices(), edge_index)
     else:
         assert torch.allclose(adj_coo.indices().flip([0]), edge_index)
@@ -165,7 +165,7 @@ def test_to_torch_csc_tensor():
     assert adj.size() == (4, 4)
     assert adj.layout == torch.sparse_csc
     adj_coo = adj.to_sparse_coo().coalesce()
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
         assert torch.allclose(adj_coo.indices(), edge_index)
         assert torch.allclose(adj_coo.values(), edge_weight)
     else:
@@ -173,7 +173,7 @@ def test_to_torch_csc_tensor():
         assert torch.allclose(adj_coo.indices()[:, perm], edge_index)
         assert torch.allclose(adj_coo.values()[perm], edge_weight)
 
-    if torch_geometric.typing.WITH_PT2:
+    if torch_geometric.typing.WITH_PT20:
         edge_attr = torch.randn(edge_index.size(1), 8)
         adj = to_torch_csc_tensor(edge_index, edge_attr)
         assert adj.size() == (4, 4, 8)