minor fixes

Author: jainapurva

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -21,14 +21,6 @@
     _linear_fp_act_fp8_weight_check,
     _linear_fp_act_fp8_weight_impl,
 )
-from torchao.prototype.dtypes.uintx.gemlite_layout import (
-    _linear_fp_act_int4_weight_gemlite_check,
-    _linear_fp_act_int4_weight_gemlite_impl,
-)
-from torchao.prototype.dtypes.floatx.floatx_tensor_core_layout import (
-    _linear_f16_bf16_act_floatx_weight_check,
-    _linear_f16_bf16_act_floatx_weight_impl,
-)
 from torchao.dtypes.uintx.int4_cpu_layout import (
     _linear_fp_act_uint4_weight_cpu_check,
     _linear_fp_act_uint4_weight_cpu_impl,

torchao/prototype/dtypes/floatx/__init__.py

@@ -4,6 +4,14 @@
 # This source code is licensed under the BSD 3-Clause license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .floatx_tensor_core_layout import FloatxTensorCoreLayout
+from .floatx_tensor_core_layout import (
+    FloatxTensorCoreLayout,
+    from_scaled_tc_floatx,
+    to_scaled_tc_floatx,
+)
 
-__all__ = ["FloatxTensorCoreLayout"]
+__all__ = [
+    "FloatxTensorCoreLayout",
+    "to_scaled_tc_floatx",
+    "from_scaled_tc_floatx",
+]