make register/deregister fn public

Author: jerryzh168

test/dtypes/test_affine_quantized.py

@@ -90,8 +90,8 @@ def test_to_device(self, apply_quant):
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     def test_register_new_dispatch(self):
         from torchao.dtypes.affine_quantized_tensor import (
-            _register_aqt_quantized_linear_dispatch,
-            _deregister_aqt_quantized_linear_dispatch,
+            register_aqt_quantized_linear_dispatch,
+            deregister_aqt_quantized_linear_dispatch,
         )
         from torchao.dtypes import to_affine_quantized_intx
         from torchao.dtypes import AffineQuantizedTensor
@@ -109,7 +109,7 @@ def impl(input_tensor, weight_tensor, bias):
             # quantized linear operator here
             assert False, "dispatching to my impl for uint6 weight only quant"
 
-        _register_aqt_quantized_linear_dispatch(dispatch_condition, impl)
+        register_aqt_quantized_linear_dispatch(dispatch_condition, impl)
 
         def apply_uint6_weight_only_quant(linear):
             linear.weight = torch.nn.Parameter(to_affine_quantized_intx(linear.weight, MappingType.ASYMMETRIC, (1, linear.weight.shape[-1]), torch.uint8, 0, 2**6-1), requires_grad=False)
@@ -122,7 +122,7 @@ def apply_uint6_weight_only_quant(linear):
         with self.assertRaisesRegex(AssertionError, "dispatching to my impl for uint6 weight only quant"):
             l(example_input)
 
-        _deregister_aqt_quantized_linear_dispatch(dispatch_condition)
+        deregister_aqt_quantized_linear_dispatch(dispatch_condition)
 
 
 

torchao/dtypes/affine_quantized_tensor.py

@@ -92,7 +92,7 @@ class QuantizedLinearNotImplementedError(NotImplementedError):
 
 
 _AQT_QLINEAR_DISPATCH_TABLE = {}
-def _register_aqt_quantized_linear_dispatch(dispatch_condition, impl):
+def register_aqt_quantized_linear_dispatch(dispatch_condition, impl):
     """Register a dispatch for quantized linear op with dispatch_condition function and impl function
     both takes three arguments:
       input_tensor: dimension is (M1, M2, ..., in_features)
@@ -108,7 +108,7 @@ def _register_aqt_quantized_linear_dispatch(dispatch_condition, impl):
     """
     _AQT_QLINEAR_DISPATCH_TABLE[dispatch_condition] = impl
 
-def _deregister_aqt_quantized_linear_dispatch(dispatch_condition):
+def deregister_aqt_quantized_linear_dispatch(dispatch_condition):
     if dispatch_condition in _AQT_QLINEAR_DISPATCH_TABLE:
         del _AQT_QLINEAR_DISPATCH_TABLE[dispatch_condition]
     else: