.

hidet-org · Jun 11, 2023 · 5a0bae1 · 5a0bae1
1 parent d675667
commit 5a0bae1
Show file tree

Hide file tree

Showing 9 changed files with 22 additions and 413 deletions.
diff --git a/experiments/matmul-parallel-k/main.py b/experiments/matmul-parallel-k/main.py
diff --git a/python/hidet/graph/ops/arithmetic.py b/python/hidet/graph/ops/arithmetic.py
@@ -49,10 +49,7 @@ def __init__(self, name: str, x: TensorNode, y: TensorNode, op: Callable[[Any, A
         inverse_map = {}
         for inp, inp_shape in zip([x, y], [x.shape, y.shape]):
             if same_list(inp_shape, z_shape):
-                inverse_map[inp] = InverseMap.from_lambda(
-                    lambda *indices: indices,
-                    num_args=len(inp_shape),
-                )
+                inverse_map[inp] = InverseMap.from_lambda(lambda *indices: indices, num_args=len(inp_shape))
             elif prod(inp_shape) == prod(z_shape):
                 inverse_map[inp] = InverseMap.from_lambda(
                     lambda *indices: [0 for _ in range(len(z_shape) - len(inp_shape))] + list(indices),