Recover pre-commit errors (#709)

Summary:
The errors were introduced by [D87842249](https://www.internalfb.com/diff/D87842249). Need to check why the diff can bypass the errors. Or is it possible to set up pre-commit runs in OSS periodically? this sounds a perfect task for AI agents.

Pull Request resolved: #709

Reviewed By: agron911

Differential Revision: D88117942

Pulled By: dshi7

fbshipit-source-id: 4753087e45bbaa500298358911599d59acb6d56a

Author: dshi7

python/test/unit/language/test_tlx.py

@@ -2311,7 +2311,7 @@ def test_stoch_round_partial_pack(dst_dtype, device):
 
     @triton.jit
     def stoch_round_partial_kernel(x_ptr, y_ptr, BLOCK_SIZE: tl.constexpr, BLOCK_SIZE_ROUNDED: tl.constexpr,
-                                    QUARTER_SIZE_ROUNDED: tl.constexpr):
+                                   QUARTER_SIZE_ROUNDED: tl.constexpr):
         # Use power-of-2 size for arange (triton requirement), then mask to actual size
         offsets_full = tl.arange(0, BLOCK_SIZE_ROUNDED)
         mask = offsets_full < BLOCK_SIZE
@@ -2357,8 +2357,8 @@ def stoch_round_partial_kernel(x_ptr, y_ptr, BLOCK_SIZE: tl.constexpr, BLOCK_SIZ
 
 
 @pytest.mark.skipif(not is_blackwell(), reason="Need Blackwell")
-@pytest.mark.parametrize("invalid_src, invalid_dst",
-                         [("float16", "float8_e5m2"), ("bfloat16", "float16"), ("float32", "int32")])
+@pytest.mark.parametrize("invalid_src, invalid_dst", [("float16", "float8_e5m2"), ("bfloat16", "float16"),
+                                                      ("float32", "int32")])
 def test_stoch_round_invalid_dtypes(invalid_src, invalid_dst, device):
     """Test that invalid dtype combinations raise proper errors."""
 
@@ -2423,6 +2423,5 @@ def stoch_round_seed_kernel(x_ptr, y_ptr, seed, BLOCK_SIZE: tl.constexpr):
 
     # Results should be different for at least some values
     different_count = (b1.float() != b2.float()).sum().item()
-    assert different_count > SIZE * 0.1, (
-        f"Different seeds should produce different results, "
-        f"but only {different_count}/{SIZE} values differ")
+    assert different_count > SIZE * 0.1, (f"Different seeds should produce different results, "
+                                          f"but only {different_count}/{SIZE} values differ")

third_party/tlx/language/tlx/utility.py

@@ -111,21 +111,16 @@ def stoch_round(
         Tensor with dtype dst_ty and shape matching src.
     """
     capability = int(cuda_parse_arch(_semantic.builder.options.arch))
-    assert capability >= 100, (
-        f"stoch_round requires compute capability >= 100 (Blackwell GPU), "
-        f"current capability: {capability}"
-    )
+    assert capability >= 100, (f"stoch_round requires compute capability >= 100 (Blackwell GPU), "
+                               f"current capability: {capability}")
     src_ty = src.type
     src_sca_ty = src_ty.scalar
 
-    assert src_sca_ty == tl.float32, (
-        f"Stochastic rounding only supports fp32 source, got {src_sca_ty}. "
-        f"Source must be float32."
-    )
-    assert dst_ty in [tl.float8e5, tl.float8e4nv, tl.float16, tl.bfloat16], (
-        f"Stochastic rounding only supports fp8/fp16/bf16 destination, got {dst_ty}. "
-        f"Supported types: float8e5 (fp8 E5M2), float8e4nv (fp8 E4M3FN), float16, bfloat16"
-    )
+    assert src_sca_ty == tl.float32, (f"Stochastic rounding only supports fp32 source, got {src_sca_ty}. "
+                                      f"Source must be float32.")
+    assert dst_ty in [tl.float8e5, tl.float8e4nv, tl.float16, tl.bfloat16
+                      ], (f"Stochastic rounding only supports fp8/fp16/bf16 destination, got {dst_ty}. "
+                          f"Supported types: float8e5 (fp8 E5M2), float8e4nv (fp8 E4M3FN), float16, bfloat16")
 
     # Verify rbits shape matches src shape
     rbits_ty = rand_bits.type