Fix flag gems tests and imports (#35)

Author: bertmaher

BackendBench/backends.py

@@ -2,9 +2,15 @@
 import importlib.util
 import logging
 from typing import Dict, Callable, List
+import torch
 
 logger = logging.getLogger(__name__)
 
+try:
+    import flag_gems
+except ImportError:
+    flag_gems = None
+
 
 class Backend:
     def __init__(self, name):
@@ -67,8 +73,6 @@ def _load_kernel_from_file(self, file_path: str, op_name: str) -> Callable:
 
     def _find_pytorch_op(self, op_name: str):
         """Map operation name to PyTorch operation."""
-        import torch
-
         # Try common patterns
         try:
             return getattr(torch.ops.aten, op_name).default
@@ -106,14 +110,10 @@ def __contains__(self, key):
 
 def _flag_gems_softmax(*args, **kwargs):
     # half_to_float is not supported in flag_gems
-    import flag_gems
-
     return flag_gems.ops.softmax(*args[:-1], **kwargs)
 
 
 def _flag_gems_layernorm(*args, **kwargs):
-    import flag_gems
-
     x, m, v = flag_gems.ops.layer_norm(*args[:-1], **kwargs)
     mv_shape = [*x.shape[:-1], 1]
     return x, m.view(*mv_shape), v.view(*mv_shape)
@@ -122,9 +122,6 @@ def _flag_gems_layernorm(*args, **kwargs):
 class FlagGemsBackend(Backend):
     def __init__(self) -> None:
         super().__init__("flaggems")
-        import flag_gems
-        import torch
-
         self.ops = {
             torch.ops.aten.abs.default: flag_gems.ops.abs,
             torch.ops.aten.abs_.default: flag_gems.ops.abs_,

test/test_backends.py

@@ -73,7 +73,7 @@ def test_flag_gems_backend_contains_op(self, mock_flag_gems):
     @patch("BackendBench.backends.flag_gems")
     def test_flag_gems_backend_getitem(self, mock_flag_gems):
         mock_abs_impl = Mock()
-        mock_flag_gems.abs = mock_abs_impl
+        mock_flag_gems.ops.abs = mock_abs_impl
 
         backend = FlagGemsBackend()