introduce mps_op_lib

Author: manuelcandales

torchchat/utils/quantize.py

@@ -137,12 +137,12 @@ def quantize_model(
                 group_size = q_kwargs["groupsize"]
                 bit_width = q_kwargs["bitwidth"]
                 has_weight_zeros = q_kwargs["has_weight_zeros"]
-                granularity = PerRow() if group_size == -1 else PerGroup(group_size) 
+                granularity = PerRow() if group_size == -1 else PerGroup(group_size)
                 weight_dtype = getattr(torch, f"int{bit_width}")
 
                 try:
                     quantize_(
-                        model, 
+                        model,
                         int8_dynamic_activation_intx_weight(
                             weight_dtype=weight_dtype,
                             granularity=granularity,
@@ -154,7 +154,7 @@ def quantize_model(
                     print("Encountered error during quantization: {e}")
                     print("Trying with PlainLayout")
                     quantize_(
-                        model, 
+                        model,
                         int8_dynamic_activation_intx_weight(
                             weight_dtype=weight_dtype,
                             granularity=granularity,
@@ -979,5 +979,19 @@ def quantized_model(self) -> nn.Module:
     except Exception as e:
         print("Unable to load torchao mps ops library.")
 
+    torchao_experimental_mps_op_lib_spec = importlib.util.spec_from_file_location(
+        "torchao_experimental_mps_op_lib",
+        f"{torchao_build_path}/src/ao/torchao/experimental/ops/mps/mps_op_lib.py",
+    )
+    torchao_experimental_mps_op_lib = importlib.util.module_from_spec(
+        torchao_experimental_mps_op_lib_spec
+    )
+    sys.modules["torchao_experimental_mps_op_lib"] = torchao_experimental_mps_op_lib
+    torchao_experimental_mps_op_lib_spec.loader.exec_module(
+        torchao_experimental_mps_op_lib
+    )
+    from torchao_experimental_mps_op_lib import *
+
+
 except Exception as e:
     print("Unable to import torchao experimental quant_api with error: ", e)