fix format check Black failed

Author: BrianPei

tests/pytorch/test_reference_activation.py

@@ -40,16 +40,12 @@
 @pytest.fixture
 def standard_input():
     # Shape (2, 4) ensures .chunk(2, dim=-1) splits it into two (2, 2) tensors cleanly
-    return torch.tensor(
-        [[-1.0, 2.0, -3.0, 4.0], [5.0, -6.0, 7.0, -8.0]], dtype=torch.float32
-    )
+    return torch.tensor([[-1.0, 2.0, -3.0, 4.0], [5.0, -6.0, 7.0, -8.0]], dtype=torch.float32)
 
 
 @pytest.fixture
 def standard_grad():
-    return torch.tensor(
-        [[0.5, 1.5, 2.5, 3.5], [4.5, 5.5, 6.5, 7.5]], dtype=torch.float32
-    )
+    return torch.tensor([[0.5, 1.5, 2.5, 3.5], [4.5, 5.5, 6.5, 7.5]], dtype=torch.float32)
 
 
 # ==============================================================================
@@ -68,9 +64,7 @@ def test_basic_forwards(standard_input):
 
     # 2. GeGLU
     a, b = standard_input.chunk(2, dim=-1)
-    assert torch.allclose(
-        geglu_torch(standard_input, quantizer), F.gelu(a, approximate="tanh") * b
-    )
+    assert torch.allclose(geglu_torch(standard_input, quantizer), F.gelu(a, approximate="tanh") * b)
 
     # 3. Quick-GeLU (qgelu)
     assert torch.allclose(
@@ -79,9 +73,7 @@ def test_basic_forwards(standard_input):
     )
 
     # 4. Quick-GeGLU (qgeglu)
-    assert torch.allclose(
-        qgeglu_torch(standard_input, quantizer), a * torch.sigmoid(1.702 * a) * b
-    )
+    assert torch.allclose(qgeglu_torch(standard_input, quantizer), a * torch.sigmoid(1.702 * a) * b)
 
     # 5. ReLU & ReGLU
     assert torch.allclose(relu_torch(standard_input, quantizer), F.relu(standard_input))
@@ -91,9 +83,7 @@ def test_basic_forwards(standard_input):
     assert torch.allclose(
         srelu_torch(standard_input, quantizer), torch.square(F.relu(standard_input))
     )
-    assert torch.allclose(
-        sreglu_torch(standard_input, quantizer), torch.square(F.relu(a)) * b
-    )
+    assert torch.allclose(sreglu_torch(standard_input, quantizer), torch.square(F.relu(a)) * b)
 
     # 7. SiLU & SwiGLU
     assert torch.allclose(silu_torch(standard_input, quantizer), F.silu(standard_input))
@@ -140,40 +130,28 @@ def test_basic_backwards(standard_grad, standard_input):
     )
 
     # 3. dqgelu & dqgeglu
-    assert (
-        dqgelu_torch(standard_grad, standard_input, quantizer).shape
-        == standard_input.shape
-    )
+    assert dqgelu_torch(standard_grad, standard_input, quantizer).shape == standard_input.shape
     assert (
         dqgeglu_torch(standard_grad[..., :2], standard_input, quantizer).shape
         == standard_input.shape
     )
 
     # 4. drelu & dreglu
-    assert (
-        drelu_torch(standard_grad, standard_input, quantizer).shape
-        == standard_input.shape
-    )
+    assert drelu_torch(standard_grad, standard_input, quantizer).shape == standard_input.shape
     assert (
         dreglu_torch(standard_grad[..., :2], standard_input, quantizer).shape
         == standard_input.shape
     )
 
     # 5. dsrelu & dsreglu
-    assert (
-        dsrelu_torch(standard_grad, standard_input, quantizer).shape
-        == standard_input.shape
-    )
+    assert dsrelu_torch(standard_grad, standard_input, quantizer).shape == standard_input.shape
     assert (
         dsreglu_torch(standard_grad[..., :2], standard_input, quantizer).shape
         == standard_input.shape
     )
 
     # 6. dsilu & dswiglu
-    assert (
-        dsilu_torch(standard_grad, standard_input, quantizer).shape
-        == standard_input.shape
-    )
+    assert dsilu_torch(standard_grad, standard_input, quantizer).shape == standard_input.shape
     assert (
         dswiglu_torch(standard_grad[..., :2], standard_input, quantizer).shape
         == standard_input.shape

tests/pytorch/test_reference_dropout.py

@@ -102,8 +102,6 @@ def test_dropout_bwd_standard_probability():
 
     # Case B: Computation directly assigned into preallocated grad_input targets
     grad_input_buffer = torch.empty_like(grad_out)
-    grad_in = dropout_bwd_torch(
-        grad_out, mask, dropout_probability=p, grad_input=grad_input_buffer
-    )
+    grad_in = dropout_bwd_torch(grad_out, mask, dropout_probability=p, grad_input=grad_input_buffer)
     assert torch.equal(grad_in, grad_input_buffer)
     assert torch.allclose(grad_input_buffer, expected_grad)

tests/pytorch/test_reference_gemm.py

@@ -53,12 +53,8 @@ def test_gemm_standard_and_device_mismatch():
     # Since out = torch.mm(B_comp, A_comp), shapes are:
     # B_comp: (M, K) = (2, 3) -> B is (2, 3) with transB=False
     # A_comp: (K, N) = (3, 2) -> A is (2, 3) with transA=True
-    A = torch.tensor(
-        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=torch.float32
-    )  # Shape (2, 3)
-    B = torch.tensor(
-        [[1.0, 0.0, 1.0], [0.0, 2.0, 1.0]], dtype=torch.float32
-    )  # Shape (2, 3)
+    A = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=torch.float32)  # Shape (2, 3)
+    B = torch.tensor([[1.0, 0.0, 1.0], [0.0, 2.0, 1.0]], dtype=torch.float32)  # Shape (2, 3)
 
     # Intentionally trigger Device Mismatch path (A on CPU, but B explicitly bound to CPU)
     # This fully exercises: if A.device != target_device: A = A.to(target_device)
@@ -152,12 +148,8 @@ def test_gemm_fp8_precision_downcast():
 
 def test_gemm_fusions_and_scaling():
     """Verify alpha scaling, bias broadcast addition, and dtype downcasting pipelines."""
-    A = torch.tensor(
-        [[2.0], [2.0]], dtype=torch.float32
-    )  # (2, 1) -> transA=False -> A_comp=(2, 1)
-    B = torch.tensor(
-        [[3.0, 4.0]], dtype=torch.float32
-    )  # (1, 2) -> transB=False -> B_comp=(1, 2)
+    A = torch.tensor([[2.0], [2.0]], dtype=torch.float32)  # (2, 1) -> transA=False -> A_comp=(2, 1)
+    B = torch.tensor([[3.0, 4.0]], dtype=torch.float32)  # (1, 2) -> transB=False -> B_comp=(1, 2)
     # torch.mm(B_comp, A_comp) -> (1, 2) x (2, 1) -> (1, 1) matrix [[14.0]]
 
     bias = torch.tensor([[1.0]], dtype=torch.float32)

transformer_engine/plugin/core/backends/reference/impl/softmax.py

@@ -33,8 +33,7 @@ def scaled_softmax_backward_torch(
     softmax_output_f32 = softmax_output.float()
 
     grad_softmax = softmax_output_f32 * (
-        output_grad_f32
-        - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
+        output_grad_f32 - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
     )
 
     return (grad_softmax * scale).to(orig_dtype)
@@ -87,8 +86,7 @@ def scaled_masked_softmax_backward_torch(
     softmax_output_f32 = softmax_output.float()
 
     grad_softmax = softmax_output_f32 * (
-        output_grad_f32
-        - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
+        output_grad_f32 - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
     )
 
     return (grad_softmax * scale).to(orig_dtype)
@@ -101,9 +99,7 @@ def scaled_upper_triang_masked_softmax_forward_torch(
     seq_len = input.size(-1)
 
     causal_mask = torch.triu(
-        torch.full(
-            (seq_len, seq_len), float("-inf"), device=input.device, dtype=input.dtype
-        ),
+        torch.full((seq_len, seq_len), float("-inf"), device=input.device, dtype=input.dtype),
         diagonal=1,
     )
 
@@ -123,8 +119,7 @@ def scaled_upper_triang_masked_softmax_backward_torch(
     softmax_output_f32 = softmax_output.float()
 
     grad_softmax = softmax_output_f32 * (
-        output_grad_f32
-        - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
+        output_grad_f32 - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
     )
 
     return (grad_softmax * scale).to(orig_dtype)
@@ -148,8 +143,7 @@ def scaled_aligned_causal_masked_softmax_backward_torch(
     softmax_output_f32 = softmax_output.float()
 
     grad_softmax = softmax_output_f32 * (
-        output_grad_f32
-        - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
+        output_grad_f32 - (softmax_output_f32 * output_grad_f32).sum(dim=-1, keepdim=True)
     )
 
     return (grad_softmax * scale).to(orig_dtype)