run black and flake8

Author: theo-barfoot

.devcontainer/devcontainer.json

@@ -23,7 +23,12 @@
                     "torchsparsegradutils/tests"
                 ],
                 "python.testing.unittestEnabled": false,
-                "python.testing.pytestEnabled": true
+                "python.testing.pytestEnabled": true,
+                "editor.formatOnSave": true,
+                "python.formatting.provider": "black",
+                "python.linting.enabled": true,
+                "python.linting.flake8Enabled": true,
+                "python.linting.pylintEnabled": false
             },
             "extensions": [
                 "dbaeumer.vscode-eslint",
@@ -34,7 +39,9 @@
                 "GitHub.vscode-pull-request-github",
                 "GitHub.vscode-github-actions",
                 "mhutchie.git-graph",
-                "waderyan.gitblame"
+                "waderyan.gitblame",
+                "ms-python.black-formatter",
+                "ms-python.flake8"
             ]
         }
     },

torchsparsegradutils/tests/benchmark_sparse_mm.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 import sys
 import os
+
 # Add the parent directory to sys.path to allow importing torchsparsegradutils
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
 import time
 import torch
@@ -19,9 +20,9 @@
 
 # problem sizes: (label, N, M, nnz)
 SIZES = [
-    ("small",  2_000,  128, 4_000),
-    ("medium", 5_000,  256, 10_000),
-    ("large",  10_000, 512, 20_000),
+    ("small", 2_000, 128, 4_000),
+    ("medium", 5_000, 256, 10_000),
+    ("large", 10_000, 512, 20_000),
 ]
 
 INDEX_DTYPES = [torch.int32, torch.int64]
@@ -33,6 +34,7 @@
     ("dense.mm", lambda A, B: torch.matmul(A.to_dense(), B)),
 ]
 
+
 def measure_op(op, A, B):
     """
     Measure forward/backward times and peak mem.
@@ -64,6 +66,7 @@ def measure_op(op, A, B):
 
     return (t1 - t0), mem_fwd, (t3 - t2), mem_bwd
 
+
 def main():
     records = []
     for size_label, N, M, nnz in SIZES:
@@ -74,10 +77,7 @@ def main():
             for val_dt in VALUE_DTYPES:
                 # build one sparse COO for all algos
                 A_coo = rand_sparse(
-                    A_shape, nnz, torch.sparse_coo,
-                    indices_dtype=idx_dt,
-                    values_dtype=val_dt,
-                    device=device
+                    A_shape, nnz, torch.sparse_coo, indices_dtype=idx_dt, values_dtype=val_dt, device=device
                 ).coalesce()
                 B = torch.randn(B_shape, dtype=val_dt, device=device)
 
@@ -89,28 +89,41 @@ def main():
                         # run
                         t_fwd, mem_fwd, t_bwd, mem_bwd = measure_op(alg_fn, A, B)
 
-                        records.append({
-                            "size":      size_label,
-                            "layout":    layout_name,
-                            "algo":      alg_name,
-                            "index_dt":  str(idx_dt).split(".")[-1],
-                            "value_dt":  str(val_dt).split(".")[-1],
-                            "N":         N,
-                            "M":         M,
-                            "nnz":       nnz,
-                            "fwd_time_s":   f"{t_fwd:.3f}",
-                            "fwd_mem_MB":   f"{mem_fwd:.1f}",
-                            "bwd_time_s":   f"{t_bwd:.3f}",
-                            "bwd_mem_MB":   f"{mem_bwd:.1f}",
-                        })
+                        records.append(
+                            {
+                                "size": size_label,
+                                "layout": layout_name,
+                                "algo": alg_name,
+                                "index_dt": str(idx_dt).split(".")[-1],
+                                "value_dt": str(val_dt).split(".")[-1],
+                                "N": N,
+                                "M": M,
+                                "nnz": nnz,
+                                "fwd_time_s": f"{t_fwd:.3f}",
+                                "fwd_mem_MB": f"{mem_fwd:.1f}",
+                                "bwd_time_s": f"{t_bwd:.3f}",
+                                "bwd_mem_MB": f"{mem_bwd:.1f}",
+                            }
+                        )
 
     df = pd.DataFrame.from_records(records)
     # reorder columns for clarity
-    df = df[[
-        "size", "layout", "algo", "index_dt", "value_dt",
-        "N", "M", "nnz",
-        "fwd_time_s", "fwd_mem_MB", "bwd_time_s", "bwd_mem_MB"
-    ]]
+    df = df[
+        [
+            "size",
+            "layout",
+            "algo",
+            "index_dt",
+            "value_dt",
+            "N",
+            "M",
+            "nnz",
+            "fwd_time_s",
+            "fwd_mem_MB",
+            "bwd_time_s",
+            "bwd_mem_MB",
+        ]
+    ]
 
     md = df.to_markdown(index=False)
     with open("torchsparsegradutils/tests/benchmark_results_sparse_mm.md", "w") as f:
@@ -120,5 +133,6 @@ def main():
 
     print("Written results to benchmark_results_sparse_mm.md")
 
+
 if __name__ == "__main__":
     main()

torchsparsegradutils/tests/profile_sparse_ops.py

@@ -3,6 +3,7 @@
 from torchsparsegradutils import sparse_mm
 from torchsparsegradutils.utils import rand_sparse
 
+
 def profile_sparse_mm(
     A_shape,
     B_shape,
@@ -11,11 +12,11 @@ def profile_sparse_mm(
     device=torch.device("cuda"),
     value_dtype=torch.float32,
     index_dtype=torch.int64,
-    row_limit: int = 20
+    row_limit: int = 20,
 ):
     """
     Profiles forward+backward of sparse_mm on a random sparse matrix.
-    
+
     Args:
         A_shape (tuple): shape of A, e.g. (N, M) or (batch, N, M)
         B_shape (tuple): shape of B, must be compatible with A: (M, P) or (batch, M, P)
@@ -30,12 +31,7 @@ def profile_sparse_mm(
         raise RuntimeError("Profiler memory breakdown requires CUDA device")
 
     # 1) build random sparse A and dense B
-    A = rand_sparse(
-        A_shape, A_nnz, layout,
-        indices_dtype=index_dtype,
-        values_dtype=value_dtype,
-        device=device
-    )
+    A = rand_sparse(A_shape, A_nnz, layout, indices_dtype=index_dtype, values_dtype=value_dtype, device=device)
     if layout == torch.sparse_coo:
         A = A.coalesce()
     B = torch.randn(*B_shape, dtype=value_dtype, device=device)
@@ -49,15 +45,12 @@ def profile_sparse_mm(
         activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
         record_shapes=True,
         profile_memory=True,
-        with_stack=False
+        with_stack=False,
     ) as prof:
         with record_function("sparse_mm_forward_and_backward"):
             out = sparse_mm(A, B)
             # you can sum or pick any scalar reduction
             out.sum().backward()
 
     # 3) print a table sorted by the CUDA memory usage of each op
-    print(prof.key_averages().table(
-        sort_by="self_cuda_memory_usage",
-        row_limit=row_limit
-    ))
+    print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=row_limit))

torchsparsegradutils/tests/test_bicgstab.py

@@ -3,16 +3,20 @@
 from torchsparsegradutils.utils import bicgstab
 
 # Device fixture
-DEVICES = [torch.device('cpu')]
+DEVICES = [torch.device("cpu")]
 if torch.cuda.is_available():
-    DEVICES.append(torch.device('cuda:0'))
+    DEVICES.append(torch.device("cuda:0"))
+
+
+def _id_device(d):
+    return str(d)
 
-def _id_device(d):   return str(d)
 
 @pytest.fixture(params=DEVICES, ids=_id_device)
 def device(request):
     return request.param
 
+
 def test_bicgstab(device):
     # setup SPD test problem
     size = 100

torchsparsegradutils/tests/test_cupy_bindings.py

@@ -7,11 +7,14 @@
 
 
 # Device fixture
-DEVICES = [torch.device('cpu')]
+DEVICES = [torch.device("cpu")]
 if torch.cuda.is_available():
-    DEVICES.append(torch.device('cuda:0'))
+    DEVICES.append(torch.device("cuda:0"))
+
+
+def _id_device(d):
+    return str(d)
 
-def _id_device(d):   return str(d)
 
 @pytest.fixture(params=DEVICES, ids=_id_device)
 def device(request):
@@ -23,24 +26,27 @@ def device(request):
     import cupy as cp
     import cupyx.scipy.sparse as csp
 
+
 # Helper to convert Cupy array to NumPy
 def _c2n(x_cupy):
     return cp.asnumpy(x_cupy) if tsgucupy.have_cupy else np.asarray(x_cupy)
 
+
 # I/O setup fixture
 @pytest.fixture
 def cupy_bindings_io(device):
     x_shape = (4, 2)
     x_t = torch.randn(x_shape, dtype=torch.float64, device=device)
     rng = np.random.default_rng()
     x_n = rng.standard_normal(x_shape, dtype=np.float64)
-    if tsgucupy.have_cupy and device.type == 'cuda':
+    if tsgucupy.have_cupy and device.type == "cuda":
         xp, xsp = cp, csp
     else:
         xp, xsp = np, nsp
     x_c = xp.asarray(x_n)
     return x_t, x_c, xsp
 
+
 # Test torch -> CuPy/NumPy COO conversion and back
 def test_t2c_and_c2t_coo(device, cupy_bindings_io):
     x_t, x_c, xsp = cupy_bindings_io
@@ -54,6 +60,7 @@ def test_t2c_and_c2t_coo(device, cupy_bindings_io):
     assert x_t2.shape == x_t_coo.shape
     assert np.allclose(x_t2.to_dense().cpu().numpy(), x_t_coo.to_dense().cpu().numpy())
 
+
 # Test NumPy/CuPy -> torch COO conversion and back
 def test_c2t_and_t2c_coo(device, cupy_bindings_io):
     x_t, x_c, xsp = cupy_bindings_io
@@ -65,6 +72,7 @@ def test_c2t_and_t2c_coo(device, cupy_bindings_io):
     assert x_c2.shape == x_c_coo.shape
     assert np.allclose(_c2n(x_c2.todense()), _c2n(x_c_coo.todense()))
 
+
 # Test torch -> CuPy/NumPy CSR conversion and back
 def test_t2c_and_c2t_csr(device, cupy_bindings_io):
     x_t, x_c, xsp = cupy_bindings_io
@@ -76,6 +84,7 @@ def test_t2c_and_c2t_csr(device, cupy_bindings_io):
     assert x_t2.shape == x_t_csr.shape
     assert np.allclose(x_t2.to_dense().cpu().numpy(), x_t_csr.to_dense().cpu().numpy())
 
+
 # Test NumPy/CuPy -> torch CSR conversion and back
 def test_c2t_and_t2c_csr(device, cupy_bindings_io):
     x_t, x_c, xsp = cupy_bindings_io

torchsparsegradutils/tests/test_cupy_sparse_solve.py

@@ -22,19 +22,24 @@
     )
 
 # Device fixture
-DEVICES = [torch.device('cpu')]
+DEVICES = [torch.device("cpu")]
 if torch.cuda.is_available():
-    DEVICES.append(torch.device('cuda:0'))
+    DEVICES.append(torch.device("cuda:0"))
+
+
+def _id_device(d):
+    return str(d)
 
-def _id_device(d):   return str(d)
 
 @pytest.fixture(params=DEVICES, ids=_id_device)
 def device(request):
     return request.param
 
+
 # relative tolerance for comparisons
 RTOL = 1e-3
 
+
 def _setup(device):
     # common setup
     A = torch.randn((4, 4), dtype=torch.float64, device=device)
@@ -44,21 +49,27 @@ def _setup(device):
     x_ref = torch.linalg.solve(A, B)
     return A_csr, B, x_ref
 
+
 def test_solver_c4t(device):
     A_csr, B, x_ref = _setup(device)
     x = tsgucupy.sparse_solve_c4t(A_csr.to(torch.float32), B.to(torch.float32))
     assert torch.allclose(x, x_ref.to(torch.float32), rtol=RTOL)
 
+
 def test_solver_gradient_c4t(device):
     A_csr, B, _ = _setup(device)
     # sparse solver gradient
-    As1 = A_csr.to(torch.float32).detach().clone(); As1.requires_grad_(True)
-    Bd1 = B.to(torch.float32).detach().clone(); Bd1.requires_grad_(True)
+    As1 = A_csr.to(torch.float32).detach().clone()
+    As1.requires_grad_(True)
+    Bd1 = B.to(torch.float32).detach().clone()
+    Bd1.requires_grad_(True)
     x = tsgucupy.sparse_solve_c4t(As1, Bd1)
     x.sum().backward()
     # dense reference gradient
-    A = As1.to_dense().detach().clone(); A.requires_grad_(True)
-    Bd2 = Bd1.detach().clone(); Bd2.requires_grad_(True)
+    A = As1.to_dense().detach().clone()
+    A.requires_grad_(True)
+    Bd2 = Bd1.detach().clone()
+    Bd2.requires_grad_(True)
     x2 = torch.linalg.solve(A, Bd2)
     x2.sum().backward()
     # compare results

torchsparsegradutils/tests/test_distributions.py

@@ -94,6 +94,7 @@ def set_seed():
 
 # Convenience functions:
 
+
 def construct_distribution(sizes, layout, var, value_dtype, index_dtype, device, requires_grad=False):
     _, batch_size, event_size, sparsity = sizes
     loc = torch.randn(event_size, device=device, dtype=value_dtype, requires_grad=requires_grad)
@@ -160,6 +161,7 @@ def check_covariance_within_tolerance(
 
 # Define Tests
 
+
 def test_rsample_forward_cov(device, layout, sizes, value_dtype, index_dtype):
     if layout == torch.sparse_coo and index_dtype == torch.int32:
         pytest.skip("Sparse COO with int32 indices is not supported")