Merge branch 'master' into trim_to_layer_pr_draft

pyg-team · Feb 28, 2023 · 1f0a3f0 · 1f0a3f0
2 parents b5b8b2d + 1760817
commit 1f0a3f0
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 37 deletions.
diff --git a/CONTRIBUTING.md → .github/CONTRIBUTING.md b/CONTRIBUTING.md → .github/CONTRIBUTING.md
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -15,10 +15,11 @@ repos:
     rev: v1.29.0
     hooks:
       - id: yamllint
-        args: [-c=.yamllint.yml]
+        name: Lint yaml
+        args: [-d, '{extends: default, rules: {line-length: disable, document-start: disable, truthy: {level: error}, braces: {max-spaces-inside: 1}}}']
 
   - repo: https://github.com/regebro/pyroma
-    rev: "4.1"
+    rev: "4.2"
     hooks:
       - id: pyroma
         name: Check packaging

diff --git a/.yamllint.yml b/.yamllint.yml
diff --git a/docs/README.md b/docs/README.md
@@ -2,7 +2,7 @@
 
 To build the documentation:
 
-1. [Build and install](https://github.com/pyg-team/pyg-lib/blob/master/CONTRIBUTING.md) `pyg-lib` from source.
+1. [Build and install](https://github.com/pyg-team/pyg-lib/blob/master/.github/CONTRIBUTING.md) `pyg-lib` from source.
 2. Install [Sphinx](https://www.sphinx-doc.org/en/master/) theme via
    ```
    pip install git+https://github.com/pyg-team/pyg_sphinx_theme.git

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -3,9 +3,9 @@
 PyG-lib Documentation
 =====================
 
-:pyg:`pyg-lib` is a **low-level Graph Neural Network library** for :pyg:`PyG` and :pytorch:`PyTorch`.
+:pyg:`null` **pyg-lib** is a **low-level Graph Neural Network library** for :pyg:`null` `PyG <https://pyg.org>`_ and :pytorch:`null` `PyTorch <https://pytorch.org>`_.
 
-.. slack::
+.. slack_button::
 
 .. toctree::
    :maxdepth: 1

diff --git a/test/ops/test_scatter_reduce.py b/test/ops/test_scatter_reduce.py
@@ -33,26 +33,37 @@ def test_fused_scatter_reduce():
     num_warmups = 1000
     num_steps = 10000
 
-    for i in range(num_warmups + num_steps):
-        if i == num_warmups:
-            torch.cuda.synchronize()
-            t = time.perf_counter()
-        out_fused = fused_scatter_reduce(x, index, dim_size=1000,
-                                         reduce_list=['sum', 'mean', 'max'])
-    torch.cuda.synchronize()
-    t = time.perf_counter() - t
-    print(f'  Fused implementation: {t:.4f} seconds')
-
-    for i in range(num_warmups + num_steps):
-        if i == num_warmups:
-            torch.cuda.synchronize()
-            t = time.perf_counter()
-        out1 = torch_scatter.scatter_add(x, index, dim_size=1000, dim=0)
-        out2 = torch_scatter.scatter_mean(x, index, dim_size=1000, dim=0)
-        out3 = torch_scatter.scatter_max(x, index, dim_size=1000, dim=0)[0]
-        out_vanilla = torch.cat([out1, out2, out3], dim=-1)
-    torch.cuda.synchronize()
-    t = time.perf_counter() - t
-    print(f'Vanilla implementation: {t:.4f} seconds')
-
-    assert torch.allclose(out_fused, out_vanilla, atol=1e-5)
+    aggrs = [
+        ['sum', 'mean'],
+        ['sum', 'mean', 'max'],
+        ['sum', 'mean', 'min', 'max'],
+    ]
+
+    for aggr in aggrs:
+        print(f'Aggregation: {aggr}')
+
+        for i in range(num_warmups + num_steps):
+            if i == num_warmups:
+                torch.cuda.synchronize()
+                t = time.perf_counter()
+            out_fused = fused_scatter_reduce(x, index, dim_size=1000,
+                                             reduce_list=aggr)
+        torch.cuda.synchronize()
+        t = time.perf_counter() - t
+        print(f'  Fused implementation: {t:.4f} seconds')
+
+        for i in range(num_warmups + num_steps):
+            if i == num_warmups:
+                torch.cuda.synchronize()
+                t = time.perf_counter()
+            outs = [
+                torch_scatter.scatter(x, index, dim_size=1000, dim=0,
+                                      reduce=reduce) for reduce in aggr
+            ]
+            out_vanilla = torch.cat(outs, dim=-1)
+        torch.cuda.synchronize()
+        t = time.perf_counter() - t
+        print(f'Vanilla implementation: {t:.4f} seconds')
+        print()
+
+        assert torch.allclose(out_fused, out_vanilla, atol=1e-5)