Repository: pyg-team/pytorch_geometric
Pull request #7216 (2023-04-20 11:51:49)

Add `batch_size` argument to `avg_pool_x` and `max_pool_x`
It can speedup runtime because:
1. We do not need to go through the batch dimension and look for max value.
2. We do not have to read tensor value which is placed on the device.

Besides dim_size can be used if a user is using fixed size datasets.

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7216?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7216](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7216?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4427254) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/e3e63d66e52aa9ca4553274f0572f1f066d99c41?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (e3e63d6) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #7216   +/-   ##
=======================================
  Coverage   91.59%   91.59%           
=======================================
  Files         437      437           
  Lines       24087    24089    +2     
=======================================
+ Hits        22062    22064    +2     
  Misses       2025     2025           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7216?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/pool/avg\_pool.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7216?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvYXZnX3Bvb2wucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/max\_pool.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7216?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvbWF4X3Bvb2wucHk=) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 560e87b60..353fb5c5a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Added an optional `batch_size` argument to `avg_pool_x` and `max_pool_x` ([#7216](https://github.com/pyg-team/pytorch_geometric/pull/7216))
 - Fixed `subgraph` on unordered inputs ([#7187](https://github.com/pyg-team/pytorch_geometric/pull/7187))
 - Allow missing node types in `HeteroDictLinear` ([#7185](https://github.com/pyg-team/pytorch_geometric/pull/7185))
 - Optimized `from_networkx` memory footprint by reducing unnecessary copies ([#7119](https://github.com/pyg-team/pytorch_geometric/pull/7119))
diff --git a/torch_geometric/nn/pool/avg_pool.py b/torch_geometric/nn/pool/avg_pool.py
index 48ea071dd..487ab4139 100644
--- a/torch_geometric/nn/pool/avg_pool.py
+++ b/torch_geometric/nn/pool/avg_pool.py
@@ -20,6 +20,7 @@ def avg_pool_x(
     cluster: Tensor,
     x: Tensor,
     batch: Tensor,
+    batch_size: Optional[int] = None,
     size: Optional[int] = None,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Average pools node features according to the clustering defined in
@@ -34,6 +35,8 @@ def avg_pool_x(
         batch (torch.Tensor): The batch vector
             :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
             node to a specific example.
+        batch_size (int, optional): The number of examples :math:`B`.
+            Automatically calculated if not given. (default: :obj:`None`)
         size (int, optional): The maximum number of clusters in a single
             example. (default: :obj:`None`)
 
@@ -41,7 +44,8 @@ def avg_pool_x(
         :obj:`None`, else :class:`torch.Tensor`
     """
     if size is not None:
-        batch_size = int(batch.max().item()) + 1
+        if batch_size is None:
+            batch_size = int(batch.max().item()) + 1
         return _avg_pool_x(cluster, x, batch_size * size), None
 
     cluster, perm = consecutive_cluster(cluster)
diff --git a/torch_geometric/nn/pool/max_pool.py b/torch_geometric/nn/pool/max_pool.py
index cf849a8cb..d099455ec 100644
--- a/torch_geometric/nn/pool/max_pool.py
+++ b/torch_geometric/nn/pool/max_pool.py
@@ -20,6 +20,7 @@ def max_pool_x(
     cluster: Tensor,
     x: Tensor,
     batch: Tensor,
+    batch_size: Optional[int] = None,
     size: Optional[int] = None,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Max-Pools node features according to the clustering defined in
@@ -33,6 +34,8 @@ def max_pool_x(
         batch (torch.Tensor): The batch vector
             :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
             node to a specific example.
+        batch_size (int, optional): The number of examples :math:`B`.
+            Automatically calculated if not given. (default: :obj:`None`)
         size (int, optional): The maximum number of clusters in a single
             example. This property is useful to obtain a batch-wise dense
             representation, *e.g.* for applying FC layers, but should only be
@@ -43,7 +46,8 @@ def max_pool_x(
         :obj:`None`, else :class:`torch.Tensor`
     """
     if size is not None:
-        batch_size = int(batch.max().item()) + 1
+        if batch_size is None:
+            batch_size = int(batch.max().item()) + 1
         return _max_pool_x(cluster, x, batch_size * size), None
 
     cluster, perm = consecutive_cluster(cluster)


Test patch:
diff --git a/test/nn/pool/test_avg_pool.py b/test/nn/pool/test_avg_pool.py
index c1e0f6dd8..15c8113df 100644
--- a/test/nn/pool/test_avg_pool.py
+++ b/test/nn/pool/test_avg_pool.py
@@ -23,11 +23,18 @@ def test_avg_pool_x():
     out, _ = avg_pool_x(cluster, x, batch, size=2)
     assert out.tolist() == [[3, 4], [5, 6], [10, 11], [0, 0]]
 
+    batch_size = int(batch.max().item()) + 1
+    out2, _ = avg_pool_x(cluster, x, batch, batch_size=batch_size, size=2)
+    assert torch.equal(out, out2)
+
     if is_full_test():
         jit = torch.jit.script(avg_pool_x)
         out, _ = jit(cluster, x, batch, size=2)
         assert out.tolist() == [[3, 4], [5, 6], [10, 11], [0, 0]]
 
+        out2, _ = jit(cluster, x, batch, batch_size=batch_size, size=2)
+        assert torch.equal(out, out2)
+
 
 def test_avg_pool():
     cluster = torch.tensor([0, 1, 0, 1, 2, 2])
diff --git a/test/nn/pool/test_max_pool.py b/test/nn/pool/test_max_pool.py
index a6c4921fb..db58de3d9 100644
--- a/test/nn/pool/test_max_pool.py
+++ b/test/nn/pool/test_max_pool.py
@@ -23,11 +23,18 @@ def test_max_pool_x():
     out, _ = max_pool_x(cluster, x, batch, size=2)
     assert out.tolist() == [[5, 6], [7, 8], [11, 12], [0, 0]]
 
+    batch_size = int(batch.max().item()) + 1
+    out2, _ = max_pool_x(cluster, x, batch, batch_size=batch_size, size=2)
+    assert torch.equal(out, out2)
+
     if is_full_test():
         jit = torch.jit.script(max_pool_x)
         out, _ = jit(cluster, x, batch, size=2)
         assert out.tolist() == [[5, 6], [7, 8], [11, 12], [0, 0]]
 
+        out2, _ = jit(cluster, x, batch, batch_size=batch_size, size=2)
+        assert torch.equal(out, out2)
+
 
 def test_max_pool():
     cluster = torch.tensor([0, 1, 0, 1, 2, 2])