Repository: pyg-team/pytorch_geometric
Pull request #9390 (2024-06-04 10:14:11)

Full repr string for `BatchNorm`
Rely on PyTorch's `extra_repr` to get more info in PyG's `BatchNorm`


Patch:
diff --git a/torch_geometric/nn/models/deep_graph_infomax.py b/torch_geometric/nn/models/deep_graph_infomax.py
index 306e905c6..e46967f96 100644
--- a/torch_geometric/nn/models/deep_graph_infomax.py
+++ b/torch_geometric/nn/models/deep_graph_infomax.py
@@ -98,7 +98,6 @@ class DeepGraphInfomax(torch.nn.Module):
         test_z: Tensor,
         test_y: Tensor,
         solver: str = 'lbfgs',
-        multi_class: str = 'auto',
         *args,
         **kwargs,
     ) -> float:
@@ -107,7 +106,7 @@ class DeepGraphInfomax(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),
diff --git a/torch_geometric/nn/models/metapath2vec.py b/torch_geometric/nn/models/metapath2vec.py
index 406d44617..f077e6f88 100644
--- a/torch_geometric/nn/models/metapath2vec.py
+++ b/torch_geometric/nn/models/metapath2vec.py
@@ -227,14 +227,13 @@ class MetaPath2Vec(torch.nn.Module):
         return pos_loss + neg_loss
 
     def test(self, train_z: Tensor, train_y: Tensor, test_z: Tensor,
-             test_y: Tensor, solver: str = "lbfgs", multi_class: str = "auto",
-             *args, **kwargs) -> float:
+             test_y: Tensor, solver: str = "lbfgs", *args, **kwargs) -> float:
         r"""Evaluates latent space quality via a logistic regression downstream
         task.
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),
diff --git a/torch_geometric/nn/models/node2vec.py b/torch_geometric/nn/models/node2vec.py
index 1b61bc5c7..e16c12191 100644
--- a/torch_geometric/nn/models/node2vec.py
+++ b/torch_geometric/nn/models/node2vec.py
@@ -173,7 +173,6 @@ class Node2Vec(torch.nn.Module):
         test_z: Tensor,
         test_y: Tensor,
         solver: str = 'lbfgs',
-        multi_class: str = 'auto',
         *args,
         **kwargs,
     ) -> float:
@@ -182,7 +181,7 @@ class Node2Vec(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),
diff --git a/torch_geometric/nn/norm/batch_norm.py b/torch_geometric/nn/norm/batch_norm.py
index 2fcc6e6ce..3830f5499 100644
--- a/torch_geometric/nn/norm/batch_norm.py
+++ b/torch_geometric/nn/norm/batch_norm.py
@@ -88,7 +88,7 @@ class BatchNorm(torch.nn.Module):
         return self.module(x)
 
     def __repr__(self):
-        return f'{self.__class__.__name__}({self.module.num_features})'
+        return f'{self.__class__.__name__}({self.module.extra_repr()})'
 
 
 class HeteroBatchNorm(torch.nn.Module):


Test patch:
diff --git a/test/nn/norm/test_batch_norm.py b/test/nn/norm/test_batch_norm.py
index 9846bc8f4..c548c0fa6 100644
--- a/test/nn/norm/test_batch_norm.py
+++ b/test/nn/norm/test_batch_norm.py
@@ -13,7 +13,8 @@ def test_batch_norm(device, conf):
     norm = BatchNorm(16, affine=conf, track_running_stats=conf).to(device)
     norm.reset_running_stats()
     norm.reset_parameters()
-    assert str(norm) == 'BatchNorm(16)'
+    assert str(norm) == (f'BatchNorm(16, eps=1e-05, momentum=0.1, '
+                         f'affine={conf}, track_running_stats={conf})')
 
     if is_full_test():
         torch.jit.script(norm)