[Misc] Black auto fix. (#4640)

* auto fix

* add more

* sort
Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>

benchmarks/benchmarks/api/bench_nn_heterographconv.py

 import time
-import dgl
-import torch
+
 import numpy as np
-import dgl.function as fn
-from dgl.nn.pytorch import SAGEConv, HeteroGraphConv
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+import dgl
+import dgl.function as fn
+from dgl.nn.pytorch import HeteroGraphConv, SAGEConv
+
 from .. import utils
 
 
-@utils.benchmark('time')
-@utils.parametrize('feat_dim', [4, 32, 256])
-@utils.parametrize('num_relations', [5, 50, 200])
+@utils.benchmark("time")
+@utils.parametrize("feat_dim", [4, 32, 256])
+@utils.parametrize("num_relations", [5, 50, 200])
 def track_time(feat_dim, num_relations):
     device = utils.get_bench_device()
     dd = {}
     nn_dict = {}
-    candidate_edges = [dgl.data.CoraGraphDataset(verbose=False)[0].edges(), dgl.data.PubmedGraphDataset(verbose=False)[
-        0].edges(), dgl.data.CiteseerGraphDataset(verbose=False)[0].edges()]
+    candidate_edges = [
+        dgl.data.CoraGraphDataset(verbose=False)[0].edges(),
+        dgl.data.PubmedGraphDataset(verbose=False)[0].edges(),
+        dgl.data.CiteseerGraphDataset(verbose=False)[0].edges(),
+    ]
     for i in range(num_relations):
-        dd[('n1', 'e_{}'.format(i), 'n2')] = candidate_edges[i %
-                                                             len(candidate_edges)]
-        nn_dict['e_{}'.format(i)] = SAGEConv(feat_dim, feat_dim, 'mean', activation=F.relu)
+        dd[("n1", "e_{}".format(i), "n2")] = candidate_edges[
+            i % len(candidate_edges)
+        ]
+        nn_dict["e_{}".format(i)] = SAGEConv(
+            feat_dim, feat_dim, "mean", activation=F.relu
+        )
 
     # dry run
     feat_dict = {}
     graph = dgl.heterograph(dd)
     for i in range(num_relations):
-        etype = 'e_{}'.format(i)
-        feat_dict[etype] = torch.randn((graph[etype].num_nodes(), feat_dim), device=device)
+        etype = "e_{}".format(i)
+        feat_dict[etype] = torch.randn(
+            (graph[etype].num_nodes(), feat_dim), device=device
+        )
 
     conv = HeteroGraphConv(nn_dict).to(device)
 

benchmarks/benchmarks/api/bench_node_subgraph.py

 import time
-import dgl
-import torch
+
 import numpy as np
-import dgl.function as fn
+import torch
 
+import dgl
+import dgl.function as fn
 
 from .. import utils
 
 
-@utils.benchmark('time')
-@utils.parametrize('graph_name', ['livejournal', 'reddit'])
-@utils.parametrize('format', ['coo', 'csc'])
-@utils.parametrize('seed_nodes_num', [200, 5000, 20000])
+@utils.benchmark("time")
+@utils.parametrize("graph_name", ["livejournal", "reddit"])
+@utils.parametrize("format", ["coo", "csc"])
+@utils.parametrize("seed_nodes_num", [200, 5000, 20000])
 def track_time(graph_name, format, seed_nodes_num):
     device = utils.get_bench_device()
     graph = utils.get_graph(graph_name, format)

benchmarks/benchmarks/api/bench_random_walk.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
+
 def _random_walk(g, seeds, length):
     return dgl.sampling.random_walk(g, seeds, length=length)
 
+
 def _node2vec(g, seeds, length):
     return dgl.sampling.node2vec_random_walk(g, seeds, 1, 1, length)
 
+
 @utils.skip_if_gpu()
-@utils.benchmark('time')
-@utils.parametrize('graph_name', ['cora', 'livejournal', 'friendster'])
-@utils.parametrize('num_seeds', [10, 100, 1000])
-@utils.parametrize('length', [2, 5, 10, 20])
-@utils.parametrize('algorithm', ['_random_walk', '_node2vec'])
+@utils.benchmark("time")
+@utils.parametrize("graph_name", ["cora", "livejournal", "friendster"])
+@utils.parametrize("num_seeds", [10, 100, 1000])
+@utils.parametrize("length", [2, 5, 10, 20])
+@utils.parametrize("algorithm", ["_random_walk", "_node2vec"])
 def track_time(graph_name, num_seeds, length, algorithm):
     device = utils.get_bench_device()
-    graph = utils.get_graph(graph_name, 'csr')
+    graph = utils.get_graph(graph_name, "csr")
     seeds = torch.randint(0, graph.num_nodes(), (num_seeds,))
     print(graph_name, num_seeds, length)
     alg = globals()[algorithm]

benchmarks/benchmarks/api/bench_readout.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
 
-@utils.benchmark('time')
-@utils.parametrize('batch_size', [4, 256, 1024])
-@utils.parametrize('feat_size', [16, 128, 512])
-@utils.parametrize('readout_op', ['sum', 'max', 'min', 'mean'])
-@utils.parametrize('type', ['edge', 'node'])
+@utils.benchmark("time")
+@utils.parametrize("batch_size", [4, 256, 1024])
+@utils.parametrize("feat_size", [16, 128, 512])
+@utils.parametrize("readout_op", ["sum", "max", "min", "mean"])
+@utils.parametrize("type", ["edge", "node"])
 def track_time(batch_size, feat_size, readout_op, type):
     device = utils.get_bench_device()
     ds = dgl.data.QM7bDataset()
@@ -17,20 +19,20 @@ def track_time(batch_size, feat_size, readout_op, type):
     graphs = ds[0:batch_size][0]
 
     g = dgl.batch(graphs).to(device)
-    if type == 'node':
-        g.ndata['h'] = torch.randn((g.num_nodes(), feat_size), device=device)
+    if type == "node":
+        g.ndata["h"] = torch.randn((g.num_nodes(), feat_size), device=device)
         for i in range(10):
-            out = dgl.readout_nodes(g, 'h', op=readout_op)
+            out = dgl.readout_nodes(g, "h", op=readout_op)
         with utils.Timer() as t:
             for i in range(50):
-                out = dgl.readout_nodes(g, 'h', op=readout_op)
-    elif type == 'edge':
-        g.edata['h'] = torch.randn((g.num_edges(), feat_size), device=device)
+                out = dgl.readout_nodes(g, "h", op=readout_op)
+    elif type == "edge":
+        g.edata["h"] = torch.randn((g.num_edges(), feat_size), device=device)
         for i in range(10):
-            out = dgl.readout_edges(g, 'h', op=readout_op)
+            out = dgl.readout_edges(g, "h", op=readout_op)
         with utils.Timer() as t:
             for i in range(50):
-                out = dgl.readout_edges(g, 'h', op=readout_op)
+                out = dgl.readout_edges(g, "h", op=readout_op)
     else:
         raise Exception("Unknown type")
 

benchmarks/benchmarks/api/bench_reverse.py

 import time
-import dgl
-import torch
+
 import numpy as np
+import torch
+
+import dgl
 
 from .. import utils
 
 
-@utils.benchmark('time', timeout=1200)
-@utils.parametrize_cpu('graph_name', ['cora', 'livejournal', 'friendster'])
-@utils.parametrize_gpu('graph_name', ['cora', 'livejournal'])
-@utils.parametrize('format', ['coo', 'csc', 'csr'])
+@utils.benchmark("time", timeout=1200)
+@utils.parametrize_cpu("graph_name", ["cora", "livejournal", "friendster"])
+@utils.parametrize_gpu("graph_name", ["cora", "livejournal"])
+@utils.parametrize("format", ["coo", "csc", "csr"])
 def track_time(graph_name, format):
     device = utils.get_bench_device()
     graph = utils.get_graph(graph_name, format)

benchmarks/benchmarks/api/bench_sample_neighbors.py

 import time
-import dgl
-import torch
+
 import numpy as np
-import dgl.function as fn
+import torch
 
+import dgl
+import dgl.function as fn
 
 from .. import utils
 
-@utils.benchmark('time')
-@utils.parametrize_cpu('graph_name', ['livejournal', 'reddit'])
-@utils.parametrize_gpu('graph_name', ['ogbn-arxiv', 'reddit'])
-@utils.parametrize('format', ['coo', 'csc'])
-@utils.parametrize('seed_nodes_num', [200, 5000, 20000])
-@utils.parametrize('fanout', [5, 20, 40])
+
+@utils.benchmark("time")
+@utils.parametrize_cpu("graph_name", ["livejournal", "reddit"])
+@utils.parametrize_gpu("graph_name", ["ogbn-arxiv", "reddit"])
+@utils.parametrize("format", ["coo", "csc"])
+@utils.parametrize("seed_nodes_num", [200, 5000, 20000])
+@utils.parametrize("fanout", [5, 20, 40])
 def track_time(graph_name, format, seed_nodes_num, fanout):
     device = utils.get_bench_device()
     graph = utils.get_graph(graph_name, format)
 
-    edge_dir = 'in'
+    edge_dir = "in"
     seed_nodes = np.random.randint(0, graph.num_nodes(), seed_nodes_num)
 
     # dry run
     for i in range(3):
         dgl.sampling.sample_neighbors(
-            graph, seed_nodes, fanout, edge_dir=edge_dir)
+            graph, seed_nodes, fanout, edge_dir=edge_dir
+        )
 
     # timing
     with utils.Timer() as t:
         for i in range(50):
             dgl.sampling.sample_neighbors(
-                graph, seed_nodes, fanout, edge_dir=edge_dir)
+                graph, seed_nodes, fanout, edge_dir=edge_dir
+            )
 
     return t.elapsed_secs / 50

benchmarks/benchmarks/api/bench_to_block.py

 import time
-import dgl
-import torch
+
 import numpy as np
+import torch
+
+import dgl
 
 from .. import utils
 
 
 @utils.skip_if_gpu()
-@utils.benchmark('time', timeout=1200)
-@utils.parametrize('graph_name', ['reddit', "ogbn-products"])
-@utils.parametrize('num_seed_nodes', [32, 256, 1024, 2048])
-@utils.parametrize('fanout', [5, 10, 20])
+@utils.benchmark("time", timeout=1200)
+@utils.parametrize("graph_name", ["reddit", "ogbn-products"])
+@utils.parametrize("num_seed_nodes", [32, 256, 1024, 2048])
+@utils.parametrize("fanout", [5, 10, 20])
 def track_time(graph_name, num_seed_nodes, fanout):
     device = utils.get_bench_device()
     data = utils.process_data(graph_name)
@@ -22,7 +24,8 @@ def track_time(graph_name, num_seed_nodes, fanout):
     subg_list = []
     for i in range(10):
         seed_nodes = np.random.randint(
-            0, graph.num_nodes(), size=num_seed_nodes)
+            0, graph.num_nodes(), size=num_seed_nodes
+        )
         subg = dgl.sampling.sample_neighbors(graph, seed_nodes, fanout)
         subg_list.append(subg)
 

benchmarks/benchmarks/api/bench_udf_apply_edges.py

 import time
-import dgl
-import torch
+
 import numpy as np
+import torch
+
+import dgl
 import dgl.function as fn
 
 from .. import utils
 
 
-@utils.benchmark('time', timeout=7200)
-@utils.parametrize('graph_name', ['ogbn-arxiv', 'pubmed'])
-@utils.parametrize('format', ['coo'])  # only coo supports udf
-@utils.parametrize('feat_size', [8, 32, 128, 512])
-@utils.parametrize('reduce_type', ['u->e', 'u+v'])
+@utils.benchmark("time", timeout=7200)
+@utils.parametrize("graph_name", ["ogbn-arxiv", "pubmed"])
+@utils.parametrize("format", ["coo"])  # only coo supports udf
+@utils.parametrize("feat_size", [8, 32, 128, 512])
+@utils.parametrize("reduce_type", ["u->e", "u+v"])
 def track_time(graph_name, format, feat_size, reduce_type):
     device = utils.get_bench_device()
     graph = utils.get_graph(graph_name, format)
     graph = graph.to(device)
-    graph.ndata['h'] = torch.randn(
-        (graph.num_nodes(), feat_size), device=device)
+    graph.ndata["h"] = torch.randn(
+        (graph.num_nodes(), feat_size), device=device
+    )
 
     reduce_udf_dict = {
-        'u->e': lambda edges: {'x': edges.src['h']},
-        'u+v': lambda edges: {'x': edges.src['h'] + edges.dst['h']},
+        "u->e": lambda edges: {"x": edges.src["h"]},
+        "u+v": lambda edges: {"x": edges.src["h"] + edges.dst["h"]},
     }
 
     # dry run

benchmarks/benchmarks/api/bench_udf_multi_update_all.py

-
 import time
-import dgl
-import torch
+
 import numpy as np
-import dgl.function as fn
+import torch
 
+import dgl
+import dgl.function as fn
 
 from .. import utils
 
 
-@utils.benchmark('time', timeout=600)
-@utils.parametrize('feat_size', [32, 128, 512])
-@utils.parametrize('num_relations', [5, 50, 500])
-@utils.parametrize('multi_reduce_type', ["sum", "stack"])
+@utils.benchmark("time", timeout=600)
+@utils.parametrize("feat_size", [32, 128, 512])
+@utils.parametrize("num_relations", [5, 50, 500])
+@utils.parametrize("multi_reduce_type", ["sum", "stack"])
 def track_time(feat_size, num_relations, multi_reduce_type):
     device = utils.get_bench_device()
     dd = {}
-    candidate_edges = [dgl.data.CoraGraphDataset(verbose=False)[0].edges(), dgl.data.PubmedGraphDataset(verbose=False)[
-        0].edges(), dgl.data.CiteseerGraphDataset(verbose=False)[0].edges()]
+    candidate_edges = [
+        dgl.data.CoraGraphDataset(verbose=False)[0].edges(),
+        dgl.data.PubmedGraphDataset(verbose=False)[0].edges(),
+        dgl.data.CiteseerGraphDataset(verbose=False)[0].edges(),
+    ]
     for i in range(num_relations):
-        dd[('n1', 'e_{}'.format(i), 'n2')] = candidate_edges[i %
-                                                             len(candidate_edges)]
+        dd[("n1", "e_{}".format(i), "n2")] = candidate_edges[
+            i % len(candidate_edges)
+        ]
     graph = dgl.heterograph(dd)
 
     graph = graph.to(device)
-    graph.nodes['n1'].data['h'] = torch.randn(
-        (graph.num_nodes('n1'), feat_size), device=device)
-    graph.nodes['n2'].data['h'] = torch.randn(
-        (graph.num_nodes('n2'), feat_size), device=device)
+    graph.nodes["n1"].data["h"] = torch.randn(
+        (graph.num_nodes("n1"), feat_size), device=device
+    )
+    graph.nodes["n2"].data["h"] = torch.randn(
+        (graph.num_nodes("n2"), feat_size), device=device
+    )
 
     # dry run
     update_dict = {}
     for i in range(num_relations):
-        update_dict['e_{}'.format(i)] = (
-            lambda edges: {'x': edges.src['h']}, lambda nodes: {'h_new': torch.sum(nodes.mailbox['x'], dim=1)})
-    graph.multi_update_all(
-        update_dict,
-        multi_reduce_type)
+        update_dict["e_{}".format(i)] = (
+            lambda edges: {"x": edges.src["h"]},
+            lambda nodes: {"h_new": torch.sum(nodes.mailbox["x"], dim=1)},
+        )
+    graph.multi_update_all(update_dict, multi_reduce_type)
 
     # timing
     with utils.Timer() as t:
         for i in range(3):
-            graph.multi_update_all(
-                update_dict,
-                multi_reduce_type)
+            graph.multi_update_all(update_dict, multi_reduce_type)
 
     return t.elapsed_secs / 3

benchmarks/benchmarks/api/bench_udf_update_all.py

 import time
-import dgl
-import torch
+
 import numpy as np
+import torch
+
+import dgl
 import dgl.function as fn
 
 from .. import utils
 
 
-@utils.benchmark('time', timeout=600)
-@utils.parametrize('graph_name', ['pubmed', 'ogbn-arxiv'])
-@utils.parametrize('format', ['coo'])  # only coo supports udf
-@utils.parametrize('feat_size', [8, 64, 512])
-@utils.parametrize('msg_type', ['copy_u', 'u_mul_e'])
-@utils.parametrize('reduce_type', ['sum', 'mean', 'max'])
+@utils.benchmark("time", timeout=600)
+@utils.parametrize("graph_name", ["pubmed", "ogbn-arxiv"])
+@utils.parametrize("format", ["coo"])  # only coo supports udf
+@utils.parametrize("feat_size", [8, 64, 512])
+@utils.parametrize("msg_type", ["copy_u", "u_mul_e"])
+@utils.parametrize("reduce_type", ["sum", "mean", "max"])
 def track_time(graph_name, format, feat_size, msg_type, reduce_type):
     device = utils.get_bench_device()
     graph = utils.get_graph(graph_name, format)
     graph = graph.to(device)
-    graph.ndata['h'] = torch.randn(
-        (graph.num_nodes(), feat_size), device=device)
-    graph.edata['e'] = torch.randn(
-        (graph.num_edges(), 1), device=device)
+    graph.ndata["h"] = torch.randn(
+        (graph.num_nodes(), feat_size), device=device
+    )
+    graph.edata["e"] = torch.randn((graph.num_edges(), 1), device=device)
 
     msg_udf_dict = {
-        'copy_u': lambda edges: {'x': edges.src['h']},
-        'u_mul_e': lambda edges: {'x': edges.src['h']*edges.data['e']},
+        "copy_u": lambda edges: {"x": edges.src["h"]},
+        "u_mul_e": lambda edges: {"x": edges.src["h"] * edges.data["e"]},
     }
 
     reduct_udf_dict = {
-        'sum': lambda nodes: {'h_new': torch.sum(nodes.mailbox['x'], dim=1)},
-        'mean': lambda nodes: {'h_new': torch.mean(nodes.mailbox['x'], dim=1)},
-        'max': lambda nodes: {'h_new': torch.max(nodes.mailbox['x'], dim=1)[0]},
+        "sum": lambda nodes: {"h_new": torch.sum(nodes.mailbox["x"], dim=1)},
+        "mean": lambda nodes: {"h_new": torch.mean(nodes.mailbox["x"], dim=1)},
+        "max": lambda nodes: {"h_new": torch.max(nodes.mailbox["x"], dim=1)[0]},
     }
 
     # dry run
@@ -39,7 +41,8 @@ def track_time(graph_name, format, feat_size, msg_type, reduce_type):
     # timing
     with utils.Timer() as t:
         for i in range(3):
-            graph.update_all(msg_udf_dict[msg_type],
-                             reduct_udf_dict[reduce_type])
+            graph.update_all(
+                msg_udf_dict[msg_type], reduct_udf_dict[reduce_type]
+            )
 
     return t.elapsed_secs / 3

benchmarks/benchmarks/api/bench_unbatch.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
-@utils.benchmark('time')
-@utils.parametrize('batch_size', [4, 32, 256, 1024])
+
+@utils.benchmark("time")
+@utils.parametrize("batch_size", [4, 32, 256, 1024])
 def track_time(batch_size):
     device = utils.get_bench_device()
     ds = dgl.data.QM7bDataset()

benchmarks/benchmarks/kernel/bench_edgesoftmax.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
+
 # The benchmarks for ops edge_softmax
-@utils.benchmark('time', timeout=600)
-@utils.parametrize('graph', ['ogbn-arxiv', 'reddit', 'cora', 'pubmed'])
-@utils.parametrize('num_heads', [1, 4, 8])
+@utils.benchmark("time", timeout=600)
+@utils.parametrize("graph", ["ogbn-arxiv", "reddit", "cora", "pubmed"])
+@utils.parametrize("num_heads", [1, 4, 8])
 def track_time(graph, num_heads):
     device = utils.get_bench_device()
     graph = utils.get_graph(graph).to(device)
-    score = torch.randn((graph.num_edges(),num_heads)).requires_grad_(True).float().to(device)
+    score = (
+        torch.randn((graph.num_edges(), num_heads))
+        .requires_grad_(True)
+        .float()
+        .to(device)
+    )
 
     # dry run
     for i in range(3):

benchmarks/benchmarks/kernel/bench_gsddmm_u_dot_v.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
+
 def calc_gflops(graph, feat_size, num_heads, time):
-    return round(2 * graph.num_edges() * feat_size / 1000000000 / time, 2)  # count both mul and add
+    return round(
+        2 * graph.num_edges() * feat_size / 1000000000 / time, 2
+    )  # count both mul and add
+
 
 # The benchmarks include broadcasting cases.
 # Given feat_size = D, num_heads = H, the node feature shape will be (H, D // H)
@@ -14,17 +20,19 @@ def calc_gflops(graph, feat_size, num_heads, time):
 #   matter how many heads are there.
 # If num_heads = 0, it falls back to the normal element-wise operation without
 #   broadcasting.
-@utils.benchmark('flops', timeout=600)
-@utils.parametrize('graph', ['ogbn-arxiv', 'reddit', 'ogbn-proteins'])
-@utils.parametrize('feat_size', [4, 32, 256])
-@utils.parametrize('num_heads', [0, 1, 4])
+@utils.benchmark("flops", timeout=600)
+@utils.parametrize("graph", ["ogbn-arxiv", "reddit", "ogbn-proteins"])
+@utils.parametrize("feat_size", [4, 32, 256])
+@utils.parametrize("num_heads", [0, 1, 4])
 def track_flops(graph, feat_size, num_heads):
     device = utils.get_bench_device()
-    graph = utils.get_graph(graph, format='coo').to(device)
+    graph = utils.get_graph(graph, format="coo").to(device)
     if num_heads == 0:
         x = torch.randn(graph.num_nodes(), feat_size, device=device)
     else:
-        x = torch.randn(graph.num_nodes(), num_heads, feat_size // num_heads, device=device)
+        x = torch.randn(
+            graph.num_nodes(), num_heads, feat_size // num_heads, device=device
+        )
 
     # dry run
     for i in range(3):

benchmarks/benchmarks/kernel/bench_gspmm_copy_u.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
+
 def calc_gflops(graph, feat_size, time):
     return round(graph.num_edges() * feat_size / 1000000000 / time, 2)
 
-@utils.benchmark('flops', timeout=600)
-@utils.parametrize('graph', ['ogbn-arxiv', 'reddit', 'ogbn-proteins'])
-@utils.parametrize('feat_size', [4, 32, 256])
-@utils.parametrize('reducer', ['sum', 'max'])
+
+@utils.benchmark("flops", timeout=600)
+@utils.parametrize("graph", ["ogbn-arxiv", "reddit", "ogbn-proteins"])
+@utils.parametrize("feat_size", [4, 32, 256])
+@utils.parametrize("reducer", ["sum", "max"])
 def track_flops(graph, feat_size, reducer):
     device = utils.get_bench_device()
-    graph = utils.get_graph(graph, format='csc').to(device)
+    graph = utils.get_graph(graph, format="csc").to(device)
     x = torch.randn(graph.num_nodes(), feat_size, device=device)
 
-    if reducer == 'sum':
+    if reducer == "sum":
         op = dgl.ops.copy_u_sum
-    elif reducer == 'max':
+    elif reducer == "max":
         op = dgl.ops.copy_u_max
     else:
-        raise ValueError('Invalid reducer', reducer)
+        raise ValueError("Invalid reducer", reducer)
 
     # dry run
     for i in range(3):

benchmarks/benchmarks/kernel/bench_gspmm_u_mul_e_sum.py

 import time
-import dgl
+
 import torch
 
+import dgl
+
 from .. import utils
 
+
 def calc_gflops(graph, feat_size, num_heads, time):
-    return round(2 * graph.num_edges() * feat_size / 1000000000 / time, 2)  # count both mul and add
+    return round(
+        2 * graph.num_edges() * feat_size / 1000000000 / time, 2
+    )  # count both mul and add
+
 
 # The benchmarks include broadcasting cases.
 # Given feat_size = D, num_heads = H, the node feature shape will be (H, D // H)
@@ -14,18 +20,20 @@ def calc_gflops(graph, feat_size, num_heads, time):
 #   matter how many heads are there.
 # If num_heads = 0, it falls back to the normal element-wise operation without
 #   broadcasting.
-@utils.benchmark('flops', timeout=600)
-@utils.parametrize('graph', ['ogbn-arxiv', 'reddit', 'ogbn-proteins'])
-@utils.parametrize('feat_size', [4, 32, 256])
-@utils.parametrize('num_heads', [0, 1, 4])
+@utils.benchmark("flops", timeout=600)
+@utils.parametrize("graph", ["ogbn-arxiv", "reddit", "ogbn-proteins"])
+@utils.parametrize("feat_size", [4, 32, 256])
+@utils.parametrize("num_heads", [0, 1, 4])
 def track_flops(graph, feat_size, num_heads):
     device = utils.get_bench_device()
-    graph = utils.get_graph(graph, format='csc').to(device)
+    graph = utils.get_graph(graph, format="csc").to(device)
     if num_heads == 0:
         x = torch.randn(graph.num_nodes(), feat_size, device=device)
         w = torch.randn(graph.num_edges(), feat_size, device=device)
     else:
-        x = torch.randn(graph.num_nodes(), num_heads, feat_size // num_heads, device=device)
+        x = torch.randn(
+            graph.num_nodes(), num_heads, feat_size // num_heads, device=device
+        )
         w = torch.randn(graph.num_edges(), num_heads, 1, device=device)
 
     # dry run

benchmarks/benchmarks/model_acc/bench_gat.py

-import dgl
-from dgl.nn.pytorch import GATConv
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import dgl
+from dgl.nn.pytorch import GATConv
+
 from .. import utils
 
+
 class GAT(nn.Module):
-    def __init__(self,
+    def __init__(
+        self,
         num_layers,
         in_dim,
         num_hidden,
@@ -17,25 +20,53 @@ class GAT(nn.Module):
         feat_drop,
         attn_drop,
         negative_slope,
-                 residual):
+        residual,
+    ):
         super(GAT, self).__init__()
         self.num_layers = num_layers
         self.gat_layers = nn.ModuleList()
         self.activation = activation
         # input projection (no residual)
-        self.gat_layers.append(GATConv(
-            in_dim, num_hidden, heads[0],
-            feat_drop, attn_drop, negative_slope, False, self.activation))
+        self.gat_layers.append(
+            GATConv(
+                in_dim,
+                num_hidden,
+                heads[0],
+                feat_drop,
+                attn_drop,
+                negative_slope,
+                False,
+                self.activation,
+            )
+        )
         # hidden layers
         for l in range(1, num_layers):
             # due to multi-head, the in_dim = num_hidden * num_heads
-            self.gat_layers.append(GATConv(
-                num_hidden * heads[l-1], num_hidden, heads[l],
-                feat_drop, attn_drop, negative_slope, residual, self.activation))
+            self.gat_layers.append(
+                GATConv(
+                    num_hidden * heads[l - 1],
+                    num_hidden,
+                    heads[l],
+                    feat_drop,
+                    attn_drop,
+                    negative_slope,
+                    residual,
+                    self.activation,
+                )
+            )
         # output projection
-        self.gat_layers.append(GATConv(
-            num_hidden * heads[-2], num_classes, heads[-1],
-            feat_drop, attn_drop, negative_slope, residual, None))
+        self.gat_layers.append(
+            GATConv(
+                num_hidden * heads[-2],
+                num_classes,
+                heads[-1],
+                feat_drop,
+                attn_drop,
+                negative_slope,
+                residual,
+                None,
+            )
+        )
 
     def forward(self, g, inputs):
         h = inputs
@@ -45,6 +76,7 @@ class GAT(nn.Module):
         logits = self.gat_layers[-1](g, h).mean(1)
         return logits
 
+
 def evaluate(model, g, features, labels, mask):
     model.eval()
     with torch.no_grad():
@@ -55,19 +87,20 @@ def evaluate(model, g, features, labels, mask):
         correct = torch.sum(indices == labels)
         return correct.item() * 1.0 / len(labels) * 100
 
-@utils.benchmark('acc')
-@utils.parametrize('data', ['cora', 'pubmed'])
+
+@utils.benchmark("acc")
+@utils.parametrize("data", ["cora", "pubmed"])
 def track_acc(data):
     data = utils.process_data(data)
     device = utils.get_bench_device()
 
     g = data[0].to(device)
 
-    features = g.ndata['feat']
-    labels = g.ndata['label']
-    train_mask = g.ndata['train_mask']
-    val_mask = g.ndata['val_mask']
-    test_mask = g.ndata['test_mask']
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
 
     in_feats = features.shape[1]
     n_classes = data.num_classes
@@ -76,17 +109,14 @@ def track_acc(data):
     g = dgl.add_self_loop(g)
 
     # create model
-    model = GAT(1, in_feats, 8, n_classes, [8, 1], F.elu,
-                0.6, 0.6, 0.2, False)
+    model = GAT(1, in_feats, 8, n_classes, [8, 1], F.elu, 0.6, 0.6, 0.2, False)
     loss_fcn = torch.nn.CrossEntropyLoss()
 
     model = model.to(device)
     model.train()
 
     # optimizer
-    optimizer = torch.optim.Adam(model.parameters(),
-                                 lr=1e-2,
-                                 weight_decay=5e-4)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
     for epoch in range(200):
         logits = model(g, features)
         loss = loss_fcn(logits[train_mask], labels[train_mask])

benchmarks/benchmarks/model_acc/bench_gcn.py

-import dgl
-from dgl.nn.pytorch import GraphConv
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import dgl
+from dgl.nn.pytorch import GraphConv
+
 from .. import utils
 
+
 class GCN(nn.Module):
-    def __init__(self,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
+    def __init__(
+        self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
         super(GCN, self).__init__()
         self.layers = nn.ModuleList()
         # input layer
         self.layers.append(GraphConv(in_feats, n_hidden, activation=activation))
         # hidden layers
         for i in range(n_layers - 1):
-            self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))
+            self.layers.append(
+                GraphConv(n_hidden, n_hidden, activation=activation)
+            )
         # output layer
         self.layers.append(GraphConv(n_hidden, n_classes))
         self.dropout = nn.Dropout(p=dropout)
@@ -33,6 +33,7 @@ class GCN(nn.Module):
             h = layer(g, h)
         return h
 
+
 def evaluate(model, g, features, labels, mask):
     model.eval()
     with torch.no_grad():
@@ -43,19 +44,20 @@ def evaluate(model, g, features, labels, mask):
         correct = torch.sum(indices == labels)
         return correct.item() * 1.0 / len(labels) * 100
 
-@utils.benchmark('acc')
-@utils.parametrize('data', ['cora', 'pubmed'])
+
+@utils.benchmark("acc")
+@utils.parametrize("data", ["cora", "pubmed"])
 def track_acc(data):
     data = utils.process_data(data)
     device = utils.get_bench_device()
 
     g = data[0].to(device).int()
 
-    features = g.ndata['feat']
-    labels = g.ndata['label']
-    train_mask = g.ndata['train_mask']
-    val_mask = g.ndata['val_mask']
-    test_mask = g.ndata['test_mask']
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
 
     in_feats = features.shape[1]
     n_classes = data.num_classes
@@ -67,7 +69,7 @@ def track_acc(data):
     degs = g.in_degrees().float()
     norm = torch.pow(degs, -0.5)
     norm[torch.isinf(norm)] = 0
-    g.ndata['norm'] = norm.unsqueeze(1)
+    g.ndata["norm"] = norm.unsqueeze(1)
 
     # create GCN model
     model = GCN(in_feats, 16, n_classes, 1, F.relu, 0.5)
@@ -77,9 +79,7 @@ def track_acc(data):
     model.train()
 
     # optimizer
-    optimizer = torch.optim.Adam(model.parameters(),
-                                 lr=1e-2,
-                                 weight_decay=5e-4)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
     for epoch in range(200):
         logits = model(g, features)
         loss = loss_fcn(logits[train_mask], labels[train_mask])

benchmarks/benchmarks/model_acc/bench_gcn_udf.py

-import dgl
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import dgl
+
 from .. import utils
 
+
 class GraphConv(nn.Module):
     def __init__(self, in_dim, out_dim, activation=None):
         super(GraphConv, self).__init__()
@@ -18,39 +20,42 @@ class GraphConv(nn.Module):
 
     def forward(self, graph, feat):
         with graph.local_scope():
-            graph.ndata['ci'] = torch.pow(graph.out_degrees().float().clamp(min=1), -0.5)
-            graph.ndata['cj'] = torch.pow(graph.in_degrees().float().clamp(min=1), -0.5)
-            graph.ndata['h'] = feat
+            graph.ndata["ci"] = torch.pow(
+                graph.out_degrees().float().clamp(min=1), -0.5
+            )
+            graph.ndata["cj"] = torch.pow(
+                graph.in_degrees().float().clamp(min=1), -0.5
+            )
+            graph.ndata["h"] = feat
             graph.update_all(self.mfunc, self.rfunc)
-            h = graph.ndata['h']
+            h = graph.ndata["h"]
             h = torch.matmul(h, self.weight) + self.bias
             if self.activation is not None:
                 h = self.activation(h)
             return h
 
     def mfunc(self, edges):
-        return {'m' : edges.src['h'], 'ci' : edges.src['ci']}
+        return {"m": edges.src["h"], "ci": edges.src["ci"]}
 
     def rfunc(self, nodes):
-        ci = nodes.mailbox['ci'].unsqueeze(2)
-        newh = (nodes.mailbox['m'] * ci).sum(1) * nodes.data['cj'].unsqueeze(1)
-        return {'h' : newh}
+        ci = nodes.mailbox["ci"].unsqueeze(2)
+        newh = (nodes.mailbox["m"] * ci).sum(1) * nodes.data["cj"].unsqueeze(1)
+        return {"h": newh}
+
 
 class GCN(nn.Module):
-    def __init__(self,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
+    def __init__(
+        self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
         super(GCN, self).__init__()
         self.layers = nn.ModuleList()
         # input layer
         self.layers.append(GraphConv(in_feats, n_hidden, activation=activation))
         # hidden layers
         for i in range(n_layers - 1):
-            self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))
+            self.layers.append(
+                GraphConv(n_hidden, n_hidden, activation=activation)
+            )
         # output layer
         self.layers.append(GraphConv(n_hidden, n_classes))
         self.dropout = nn.Dropout(p=dropout)
@@ -63,6 +68,7 @@ class GCN(nn.Module):
             h = layer(g, h)
         return h
 
+
 def evaluate(model, g, features, labels, mask):
     model.eval()
     with torch.no_grad():
@@ -73,19 +79,20 @@ def evaluate(model, g, features, labels, mask):
         correct = torch.sum(indices == labels)
         return correct.item() * 1.0 / len(labels) * 100
 
-@utils.benchmark('acc', timeout=300)
-@utils.parametrize('data', ['cora', 'pubmed'])
+
+@utils.benchmark("acc", timeout=300)
+@utils.parametrize("data", ["cora", "pubmed"])
 def track_acc(data):
     data = utils.process_data(data)
     device = utils.get_bench_device()
 
     g = data[0].to(device).int()
 
-    features = g.ndata['feat']
-    labels = g.ndata['label']
-    train_mask = g.ndata['train_mask']
-    val_mask = g.ndata['val_mask']
-    test_mask = g.ndata['test_mask']
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
 
     in_feats = features.shape[1]
     n_classes = data.num_classes
@@ -97,7 +104,7 @@ def track_acc(data):
     degs = g.in_degrees().float()
     norm = torch.pow(degs, -0.5)
     norm[torch.isinf(norm)] = 0
-    g.ndata['norm'] = norm.unsqueeze(1)
+    g.ndata["norm"] = norm.unsqueeze(1)
 
     # create GCN model
     model = GCN(in_feats, 16, n_classes, 1, F.relu, 0.5)
@@ -107,9 +114,7 @@ def track_acc(data):
     model.train()
 
     # optimizer
-    optimizer = torch.optim.Adam(model.parameters(),
-                                 lr=1e-2,
-                                 weight_decay=5e-4)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
     for epoch in range(200):
         logits = model(g, features)
         loss = loss_fcn(logits[train_mask], labels[train_mask])

benchmarks/benchmarks/model_acc/bench_rgcn_base.py

@@ -2,41 +2,51 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchmetrics.functional import accuracy
-from .. import utils
-from .. import rgcn
 
+from .. import rgcn, utils
 
-@utils.benchmark('acc', timeout=1200)
-@utils.parametrize('dataset', ['aifb', 'mutag'])
-@utils.parametrize('ns_mode', [False])
+
+@utils.benchmark("acc", timeout=1200)
+@utils.parametrize("dataset", ["aifb", "mutag"])
+@utils.parametrize("ns_mode", [False])
 def track_acc(dataset, ns_mode):
-    g, num_rels, num_classes, labels, train_idx, test_idx, target_idx = rgcn.load_data(
-        dataset, get_norm=True)
+    (
+        g,
+        num_rels,
+        num_classes,
+        labels,
+        train_idx,
+        test_idx,
+        target_idx,
+    ) = rgcn.load_data(dataset, get_norm=True)
     num_hidden = 16
-    if dataset == 'aifb':
+    if dataset == "aifb":
         num_bases = -1
-        l2norm = 0.
-    elif dataset == 'mutag':
+        l2norm = 0.0
+    elif dataset == "mutag":
         num_bases = 30
         l2norm = 5e-4
-    elif dataset == 'am':
+    elif dataset == "am":
         num_bases = 40
         l2norm = 5e-4
     else:
         raise ValueError()
-    model = rgcn.RGCN(g.num_nodes(),
+    model = rgcn.RGCN(
+        g.num_nodes(),
         num_hidden,
         num_classes,
         num_rels,
         num_bases=num_bases,
-                 ns_mode=ns_mode)
+        ns_mode=ns_mode,
+    )
     device = utils.get_bench_device()
     labels = labels.to(device)
     model = model.to(device)
     g = g.int().to(device)
 
     optimizer = torch.optim.Adam(
-        model.parameters(), lr=1e-2, weight_decay=l2norm)
+        model.parameters(), lr=1e-2, weight_decay=l2norm
+    )
 
     model.train()
     for epoch in range(30):
@@ -51,7 +61,6 @@ def track_acc(dataset, ns_mode):
     with torch.no_grad():
         logits = model(g)
     logits = logits[target_idx]
-    test_acc = accuracy(logits[test_idx].argmax(
-        dim=1), labels[test_idx]).item()
+    test_acc = accuracy(logits[test_idx].argmax(dim=1), labels[test_idx]).item()
 
     return test_acc

benchmarks/benchmarks/model_acc/bench_sage.py

-import dgl
-from dgl.nn.pytorch import SAGEConv
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import dgl
+from dgl.nn.pytorch import SAGEConv
+
 from .. import utils
 
+
 class GraphSAGE(nn.Module):
-    def __init__(self,
+    def __init__(
+        self,
         in_feats,
         n_hidden,
         n_classes,
         n_layers,
         activation,
         dropout,
-                 aggregator_type):
+        aggregator_type,
+    ):
         super(GraphSAGE, self).__init__()
         self.layers = nn.ModuleList()
         self.dropout = nn.Dropout(dropout)
@@ -26,7 +30,9 @@ class GraphSAGE(nn.Module):
         for i in range(n_layers - 1):
             self.layers.append(SAGEConv(n_hidden, n_hidden, aggregator_type))
         # output layer
-        self.layers.append(SAGEConv(n_hidden, n_classes, aggregator_type)) # activation None
+        self.layers.append(
+            SAGEConv(n_hidden, n_classes, aggregator_type)
+        )  # activation None
 
     def forward(self, graph, inputs):
         h = self.dropout(inputs)
@@ -37,6 +43,7 @@ class GraphSAGE(nn.Module):
                 h = self.dropout(h)
         return h
 
+
 def evaluate(model, g, features, labels, mask):
     model.eval()
     with torch.no_grad():
@@ -47,19 +54,20 @@ def evaluate(model, g, features, labels, mask):
         correct = torch.sum(indices == labels)
         return correct.item() * 1.0 / len(labels) * 100
 
-@utils.benchmark('acc')
-@utils.parametrize('data', ['cora', 'pubmed'])
+
+@utils.benchmark("acc")
+@utils.parametrize("data", ["cora", "pubmed"])
 def track_acc(data):
     data = utils.process_data(data)
     device = utils.get_bench_device()
 
     g = data[0].to(device)
 
-    features = g.ndata['feat']
-    labels = g.ndata['label']
-    train_mask = g.ndata['train_mask']
-    val_mask = g.ndata['val_mask']
-    test_mask = g.ndata['test_mask']
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
 
     in_feats = features.shape[1]
     n_classes = data.num_classes
@@ -68,16 +76,14 @@ def track_acc(data):
     g = dgl.add_self_loop(g)
 
     # create model
-    model = GraphSAGE(in_feats, 16, n_classes, 1, F.relu, 0.5, 'gcn')
+    model = GraphSAGE(in_feats, 16, n_classes, 1, F.relu, 0.5, "gcn")
     loss_fcn = torch.nn.CrossEntropyLoss()
 
     model = model.to(device)
     model.train()
 
     # optimizer
-    optimizer = torch.optim.Adam(model.parameters(),
-                                 lr=1e-2,
-                                 weight_decay=5e-4)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
     for epoch in range(200):
         logits = model(g, features)
         loss = loss_fcn(logits[train_mask], labels[train_mask])