ran black and isort for coherent code formatting

1 files changed, 137 insertions, 63 deletions

diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index d2b5d59..9cd01a9 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -1,27 +1,38 @@
-import pytest
+from itertools import product
 
+import pytest
 import torch
-import bitsandbytes as bnb
 
-from itertools import product
+import bitsandbytes as bnb
 
 n = 1
 k = 25
-dim1 = torch.randint(16,64, size=(n,)).tolist()
-dim2 = torch.randint(32,96, size=(n,)).tolist()
-dim3 = torch.randint(32,96, size=(n,)).tolist()
-dim4 = torch.randint(32,96, size=(n,)).tolist()
+dim1 = torch.randint(16, 64, size=(n,)).tolist()
+dim2 = torch.randint(32, 96, size=(n,)).tolist()
+dim3 = torch.randint(32, 96, size=(n,)).tolist()
+dim4 = torch.randint(32, 96, size=(n,)).tolist()
 funcs = [(torch.bmm, bnb.bmm_cublas), (torch.matmul, bnb.matmul_cublas)]
-str_funcs = ['bmm', 'matmul']
+str_funcs = ["bmm", "matmul"]
 req_grad = [(False, False), (True, False), (True, True), (False, True)]
-req_grad_str = ['FF', 'TF', 'TT', 'FT']
+req_grad_str = ["FF", "TF", "TT", "FT"]
 transpose = [(False, False), (False, True), (True, True), (True, False)]
-str_transpose = ['FF', 'FT', 'TT', 'TF']
+str_transpose = ["FF", "FT", "TT", "TF"]
 dtype = [torch.float32, torch.float16]
-values = list(product(dim1,dim2,dim3,dim4,funcs, dtype, req_grad, transpose))
-str_values = list(product(dim1,dim2,dim3,dim4,str_funcs, dtype, req_grad_str, str_transpose))
-names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_func_{4}_dtype_{5}_requires_grad_{6}_transpose_{7}'.format(*vals) for vals in str_values]
-@pytest.mark.parametrize("dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose", values, ids=names)
+values = list(product(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose))
+str_values = list(
+    product(dim1, dim2, dim3, dim4, str_funcs, dtype, req_grad_str, str_transpose)
+)
+names = [
+    "dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_func_{4}_dtype_{5}_requires_grad_{6}_transpose_{7}".format(
+        *vals
+    )
+    for vals in str_values
+]
+
+
+@pytest.mark.parametrize(
+    "dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose", values, ids=names
+)
 def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
     dim2 = dim2 - (dim2 % 16)
     dim3 = dim3 - (dim3 % 16)
@@ -32,9 +43,11 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
         if funcs[0] in [torch.mm, torch.matmul]:
             dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2)
             dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3)
-            A = torch.randn(size=dimA, device='cuda', requires_grad=req_grad[0])
-            B = torch.randn(size=dimB, device='cuda', requires_grad=req_grad[1])
-            target = torch.randn(size=(dim2, dim4), device='cuda', requires_grad=req_grad[1])
+            A = torch.randn(size=dimA, device="cuda", requires_grad=req_grad[0])
+            B = torch.randn(size=dimB, device="cuda", requires_grad=req_grad[1])
+            target = torch.randn(
+                size=(dim2, dim4), device="cuda", requires_grad=req_grad[1]
+            )
             torch.nn.init.xavier_uniform_(B)
 
             if not transpose[0] and not transpose[1]:
@@ -52,9 +65,9 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
 
             n = out_bnb.numel()
             idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
-            assert (idx==0).sum().item() < n*0.0175
+            assert (idx == 0).sum().item() < n * 0.0175
             idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2)
-            assert (idx==0).sum().item() < n*0.001
+            assert (idx == 0).sum().item() < n * 0.001
 
             if any(req_grad):
                 out_bnb.data.copy_(out_torch)
@@ -78,16 +91,22 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
             if req_grad[1]:
                 n = gradB1.numel()
                 idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
-                assert (idx==0).sum().item() < n*0.1
+                assert (idx == 0).sum().item() < n * 0.1
                 idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
-                assert (idx==0).sum().item() < n*0.02
+                assert (idx == 0).sum().item() < n * 0.02
                 torch.testing.assert_allclose(gradB1, gradB2, atol=0.18, rtol=0.3)
 
         # batched matrix multiply
         if funcs[0] in [torch.bmm, torch.matmul]:
-            A = torch.randn(size=(dim1, dim2, dim3), device='cuda', requires_grad=req_grad[0])
-            B = torch.randn(size=(dim1, dim3, dim4), device='cuda', requires_grad=req_grad[1])
-            target = torch.randn(size=(dim1, dim2, dim4), device='cuda', requires_grad=req_grad[1])
+            A = torch.randn(
+                size=(dim1, dim2, dim3), device="cuda", requires_grad=req_grad[0]
+            )
+            B = torch.randn(
+                size=(dim1, dim3, dim4), device="cuda", requires_grad=req_grad[1]
+            )
+            target = torch.randn(
+                size=(dim1, dim2, dim4), device="cuda", requires_grad=req_grad[1]
+            )
             torch.nn.init.xavier_uniform_(B)
 
             out_torch = funcs[0](A, B)
@@ -95,7 +114,7 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
 
             n = out_bnb.numel()
             idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
-            assert (idx==0).sum().item() < n*0.01
+            assert (idx == 0).sum().item() < n * 0.01
             torch.testing.assert_allclose(out_bnb, out_torch, atol=0.027, rtol=0.2)
 
             if any(req_grad):
@@ -120,16 +139,20 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
             if req_grad[1]:
                 n = gradB1.numel()
                 idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
-                assert (idx==0).sum().item() < n*0.1
+                assert (idx == 0).sum().item() < n * 0.1
                 idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
-                assert (idx==0).sum().item() < n*0.02
+                assert (idx == 0).sum().item() < n * 0.02
 
         if funcs[0] in [torch.matmul]:
             dim1 = dim1 - (dim1 % 16)
-            A = torch.randn(size=(dim1, dim2, dim3), device='cuda', requires_grad=req_grad[0])
+            A = torch.randn(
+                size=(dim1, dim2, dim3), device="cuda", requires_grad=req_grad[0]
+            )
             dimB = (dim4, dim3) if transpose[1] else (dim3, dim4)
-            B = torch.randn(size=dimB, device='cuda', requires_grad=req_grad[1])
-            target = torch.randn(size=(dim1, dim2, dim4), device='cuda', requires_grad=req_grad[1])
+            B = torch.randn(size=dimB, device="cuda", requires_grad=req_grad[1])
+            target = torch.randn(
+                size=(dim1, dim2, dim4), device="cuda", requires_grad=req_grad[1]
+            )
             torch.nn.init.xavier_uniform_(B)
 
             if transpose[1]:
@@ -141,9 +164,9 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
 
             n = out_bnb.numel()
             idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
-            assert (idx==0).sum().item() < n*0.0175
+            assert (idx == 0).sum().item() < n * 0.0175
             idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2)
-            assert (idx==0).sum().item() < n*0.001
+            assert (idx == 0).sum().item() < n * 0.001
 
             if any(req_grad):
                 out_bnb.data.copy_(out_torch)
@@ -167,51 +190,96 @@ def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
             if req_grad[1]:
                 n = gradB1.numel()
                 idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
-                assert (idx==0).sum().item() < n*0.1
+                assert (idx == 0).sum().item() < n * 0.1
                 idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
-                assert (idx==0).sum().item() < n*0.02
+                assert (idx == 0).sum().item() < n * 0.02
 
 
 n = 1
 k = 3
-dim1 = torch.randint(16,64, size=(n,)).tolist()
-dim2 = torch.randint(32,96, size=(n,)).tolist()
-dim3 = torch.randint(32,96, size=(n,)).tolist()
-dim4 = torch.randint(32,96, size=(n,)).tolist()
+dim1 = torch.randint(16, 64, size=(n,)).tolist()
+dim2 = torch.randint(32, 96, size=(n,)).tolist()
+dim3 = torch.randint(32, 96, size=(n,)).tolist()
+dim4 = torch.randint(32, 96, size=(n,)).tolist()
 
-#dim1 = (17,)
-#dim2 = (7,)
-#dim3 = (37,)
-#dim4 = (23,)
+# dim1 = (17,)
+# dim2 = (7,)
+# dim3 = (37,)
+# dim4 = (23,)
 
 decomp = [0.0, 6.0]
 funcs = [(torch.matmul, bnb.matmul)]
-str_funcs = ['matmul']
+str_funcs = ["matmul"]
 req_grad = [(False, False), (True, False), (True, True), (False, True)]
-req_grad_str = ['FF', 'TF', 'TT', 'FT']
+req_grad_str = ["FF", "TF", "TT", "FT"]
 transpose = [(False, True), (False, False)]
-str_transpose = ['NT', 'NN']
+str_transpose = ["NT", "NN"]
 dtype = [torch.float16]
 has_fp16_weights = [True, False]
-values = list(product(dim1,dim2,dim3,dim4,funcs, dtype, req_grad, transpose, decomp, has_fp16_weights))
-str_values = list(product(dim1,dim2,dim3,dim4,str_funcs, dtype, req_grad_str, str_transpose, decomp, has_fp16_weights))
-names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_func_{4}_dtype_{5}_requires_grad_{6}_transpose_{7}_decomp_{8}_has_fp16_weights_{9}'.format(*vals) for vals in str_values]
-@pytest.mark.parametrize("dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights", values, ids=names)
-def test_matmullt(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights):
+values = list(
+    product(
+        dim1,
+        dim2,
+        dim3,
+        dim4,
+        funcs,
+        dtype,
+        req_grad,
+        transpose,
+        decomp,
+        has_fp16_weights,
+    )
+)
+str_values = list(
+    product(
+        dim1,
+        dim2,
+        dim3,
+        dim4,
+        str_funcs,
+        dtype,
+        req_grad_str,
+        str_transpose,
+        decomp,
+        has_fp16_weights,
+    )
+)
+names = [
+    "dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_func_{4}_dtype_{5}_requires_grad_{6}_transpose_{7}_decomp_{8}_has_fp16_weights_{9}".format(
+        *vals
+    )
+    for vals in str_values
+]
+
+
+@pytest.mark.parametrize(
+    "dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights",
+    values,
+    ids=names,
+)
+def test_matmullt(
+    dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights
+):
     dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2)
     dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3)
-    outlier_dim = torch.randint(0, dimA[1], size=(dimA[1]//8,), device='cuda')
+    outlier_dim = torch.randint(0, dimA[1], size=(dimA[1] // 8,), device="cuda")
 
     for i in range(k):
 
         # normal multiply
         if funcs[0] in [torch.mm, torch.matmul]:
-            A = torch.randn(size=dimA, device='cuda', requires_grad=req_grad[0], dtype=dtype)
+            A = torch.randn(
+                size=dimA, device="cuda", requires_grad=req_grad[0], dtype=dtype
+            )
             if decomp == 6.0:
                 with torch.no_grad():
                     A[:, outlier_dim] = 6.0
-            B = torch.randn(size=dimB, device='cuda', requires_grad=req_grad[1], dtype=dtype)
-            target = torch.randn(size=(dim2, dim4), device='cuda', requires_grad=req_grad[1], dtype=dtype)
+            B = torch.randn(
+                size=dimB, device="cuda", requires_grad=req_grad[1], dtype=dtype
+            )
+            target = torch.randn(
+                size=(dim2, dim4), device="cuda", requires_grad=req_grad[1], dtype=dtype
+            )
             torch.nn.init.xavier_uniform_(B)
             B2 = B.clone()
 
@@ -219,8 +287,15 @@ def test_matmullt(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, dec
             state.threshold = decomp
             state.has_fp16_weights = has_fp16_weights
             if not has_fp16_weights:
-                if not transpose[0] and not transpose[1]: B2 = B2.t().contiguous()
-                state.CB, CBt, state.SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B2)
+                if not transpose[0] and not transpose[1]:
+                    B2 = B2.t().contiguous()
+                (
+                    state.CB,
+                    CBt,
+                    state.SCB,
+                    SCBt,
+                    coo_tensorB,
+                ) = bnb.functional.double_quant(B2)
                 B2 = state.CB
 
             if not transpose[0] and transpose[1]:
@@ -231,12 +306,12 @@ def test_matmullt(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, dec
                 out_bnb = funcs[1](A, B2.t(), state=state)
 
             n = out_bnb.numel()
-            err = torch.abs(out_bnb-out_torch).mean().item()
-            #print(f'abs error {err:.4f}')
+            err = torch.abs(out_bnb - out_torch).mean().item()
+            # print(f'abs error {err:.4f}')
             idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
-            assert (idx==0).sum().item() < n*0.0175
+            assert (idx == 0).sum().item() < n * 0.0175
             idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2)
-            assert (idx==0).sum().item() < n*0.001
+            assert (idx == 0).sum().item() < n * 0.001
 
             if has_fp16_weights:
                 if any(req_grad):
@@ -263,8 +338,7 @@ def test_matmullt(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, dec
                     assert torch.abs(gradB1).sum() > 0.0
                     assert torch.abs(gradB2).sum() > 0.0
                     idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
-                    assert (idx==0).sum().item() < n*0.1
+                    assert (idx == 0).sum().item() < n * 0.1
                     idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
-                    assert (idx==0).sum().item() < n*0.02
+                    assert (idx == 0).sum().item() < n * 0.02
                     torch.testing.assert_allclose(gradB1, gradB2, atol=0.18, rtol=0.3)
-