format change

Signed-off-by: gdeng <gdeng@nvidia.com>

Author: gdengk

emerging_optimizers/triton_kernels/syrk.py

@@ -17,6 +17,7 @@
 import triton
 import triton.language as tl
 
+
 try:
     from triton.tools.tensor_descriptor import TensorDescriptor
 except ImportError:
@@ -322,9 +323,9 @@ def tsyrk_ex(
     assert a.is_contiguous() or a.T.is_contiguous(), "invalid input tensor layout. a or a.T must be contiguous."
 
     N, K = a.shape
-    assert (c is None and beta == 0.0) or (
-        c is not None and c.shape == (N, N)
-    ), "if c is provided, c must be of shape (N, N)"
+    assert (c is None and beta == 0.0) or (c is not None and c.shape == (N, N)), (
+        "if c is provided, c must be of shape (N, N)"
+    )
     assert c is None or c.is_contiguous() or c.T.is_contiguous(), "if c is provided, c or c.T must be contiguous"
 
     d = torch.empty((N, N), device=a.device, dtype=a.dtype)

tests/test_triton_kernels.py

@@ -12,15 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from absl.testing import absltest, parameterized
 import torch
+from absl.testing import absltest, parameterized
 
-from emerging_optimizers import triton_kernels
-from emerging_optimizers import utils
+from emerging_optimizers import triton_kernels, utils
 
 
 class TritonKernelsTest(parameterized.TestCase):
-
     @parameterized.product(
         ({"n": 5, "k": 7, "rtol": 1e-5}, {"n": 17, "k": 23, "rtol": 1e-3}, {"n": 127, "k": 255, "rtol": 2e-2}),
         ({"trans": False},),
@@ -52,7 +50,6 @@ def test_ssyrk_tf32_not_far_from_matmul(self, n: int, k: int, trans: bool, rtol:
 
 
 class TritonKernelsIntegerInputTest(parameterized.TestCase):
-
     @parameterized.product(
         ({"n": 5, "k": 7}, {"n": 17, "k": 23}, {"n": 127, "k": 255}),
         ({"trans": False},),