refactor(coreai_opt): apply internal-import _-alias fixes

- Add the required `_` alias to private-module imports in public modules, drop the redundant `_` alias in private modules, and rename references to match, per code style guide §3.3
- Generated with `python scripts/pre_commit/check_internal_import_aliases.py --fix`
- Kept as a separate commit from the hook so a future rebase conflict here can be resolved by reverting `src/` and re-running `--fix`, instead of hand-merging auto-generated edits

Author: dengqiaoyu

src/coreai_opt/_utils/insertion/torch_function/modes.py

@@ -13,7 +13,7 @@
 import torch
 import torch.nn as nn
 
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt._utils.torch_utils import NamedModule
 from coreai_opt.config.spec import CompressionSimulatorBase
 
@@ -129,7 +129,7 @@ def resolve(
         func: Callable,
         state_tensor: torch.Tensor,
         current_module: NamedModule,
-        components_dict: Mapping[int | str, _PartialConstructor | None],
+        components_dict: Mapping[int | str, PartialConstructor | None],
     ) -> None:
         """Resolve and cache the optimizer for ``state_tensor`` at the current call site.
 
@@ -197,7 +197,7 @@ def _resolve_op_state(
         self,
         func: Callable,
         state_tensor: torch.Tensor,
-        components_dict: Mapping[int | str, _PartialConstructor | None],
+        components_dict: Mapping[int | str, PartialConstructor | None],
     ) -> CompressionSimulatorBase | None:
         """Look up the ``op_state_spec`` optimizer using all local names for ``state_tensor``."""
         local_state_names = self.get_all_local_names(state_tensor)

src/coreai_opt/_utils/insertion/torch_function/state_spec_resolver.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt.config.spec import CompressionSimulatorBase
 
 
@@ -27,15 +27,15 @@ class OpCompressionComponents:
     - None if no compression is applied to that component
     """
 
-    op_input_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
-    ] = field(default_factory=dict)
+    op_input_components: Mapping[int | str, PartialConstructor[CompressionSimulatorBase] | None] = (
+        field(default_factory=dict)
+    )
 
     op_output_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
+        int | str, PartialConstructor[CompressionSimulatorBase] | None
     ] = field(default_factory=dict)
 
-    op_state_components: Mapping[str, _PartialConstructor[CompressionSimulatorBase] | None] = field(
+    op_state_components: Mapping[str, PartialConstructor[CompressionSimulatorBase] | None] = field(
         default_factory=dict
     )
 
@@ -74,15 +74,15 @@ class ModuleCompressionComponents:
     or name to an OpCompressionComponents class.
     """
 
-    weight: Mapping[str, _PartialConstructor[CompressionSimulatorBase] | None] = field(
+    weight: Mapping[str, PartialConstructor[CompressionSimulatorBase] | None] = field(
         default_factory=dict
     )
 
-    input_activation: Mapping[int | str, _PartialConstructor[CompressionSimulatorBase] | None] = (
+    input_activation: Mapping[int | str, PartialConstructor[CompressionSimulatorBase] | None] = (
         field(default_factory=dict)
     )
 
-    output_activation: Mapping[int | str, _PartialConstructor[CompressionSimulatorBase] | None] = (
+    output_activation: Mapping[int | str, PartialConstructor[CompressionSimulatorBase] | None] = (
         field(default_factory=dict)
     )
 
@@ -91,14 +91,14 @@ class ModuleCompressionComponents:
     op_name_components: Mapping[str, OpCompressionComponents] = field(default_factory=dict)
 
     module_input_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
+        int | str, PartialConstructor[CompressionSimulatorBase] | None
     ] = field(default_factory=dict)
 
     module_output_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
+        int | str, PartialConstructor[CompressionSimulatorBase] | None
     ] = field(default_factory=dict)
 
-    module_state_components: Mapping[str, _PartialConstructor[CompressionSimulatorBase] | None] = (
+    module_state_components: Mapping[str, PartialConstructor[CompressionSimulatorBase] | None] = (
         field(default_factory=dict)
     )
 

src/coreai_opt/_utils/insertion/torch_function/types.py

@@ -16,7 +16,7 @@
 from torch.fx.operator_schemas import create_type_hint, normalize_function
 
 from coreai_opt._utils.config_utils import get_last_matching_spec
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt.config.spec import CompressionSimulatorBase
 
 logger = logging.getLogger(__name__)
@@ -130,7 +130,7 @@ def any_tensor_optimizable(args: list[Any], kwargs: dict[str, Any]) -> bool:
 def get_optimizer_from_components_dict(
     func: Callable,
     tensor_identifiers: int | str | list[int | str],
-    components_dict: Mapping[int | str, _PartialConstructor | None],
+    components_dict: Mapping[int | str, PartialConstructor | None],
 ) -> tuple[CompressionSimulatorBase | None, bool]:
     """Return the appropriate optimizer from ``components_dict``.
 

src/coreai_opt/_utils/insertion/torch_function/utils.py

@@ -17,7 +17,7 @@
 import torch
 import torch.nn.utils.parametrize as P
 
-from coreai_opt._utils.version_utils import version_ge as _version_ge
+from coreai_opt._utils.version_utils import version_ge
 
 logger = logging.getLogger(__name__)
 
@@ -424,7 +424,7 @@ def export_model(
             exported_program = torch.export.export(
                 model, example_inputs, dynamic_shapes=dynamic_shapes
             )
-            if _version_ge(torch, "2.9"):
+            if version_ge(torch, "2.9"):
                 exported_model = exported_program.module(check_guards=False)
             else:
                 exported_model = exported_program.module()

src/coreai_opt/_utils/torch_utils.py

@@ -10,10 +10,10 @@
 import torch
 import torch.nn as nn
 
-from coreai_opt._utils.registry_utils import ClassRegistryMixin
+from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin
 
 
-class CompressionSimulatorBase(ClassRegistryMixin, nn.Module):
+class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     """
     Abstract base class for compression simulators.
 

src/coreai_opt/config/spec/compression_simulator.py

@@ -14,17 +14,17 @@
 import numpy as np
 
 from coreai_opt.coreai_utils._coreai_imports import (
-    AIProgram,
-    DenseResourceElementsAttr,
-    F16Type,
-    F32Type,
-    InsertionPoint,
-    IntegerType,
-    RankedTensorType,
-    WalkResult,
+    AIProgram as _AIProgram,
+    DenseResourceElementsAttr as _DenseResourceElementsAttr,
+    F16Type as _F16Type,
+    F32Type as _F32Type,
+    InsertionPoint as _InsertionPoint,
+    IntegerType as _IntegerType,
+    RankedTensorType as _RankedTensorType,
+    WalkResult as _WalkResult,
     _get_constant_value_as_np_array,
-    compression_types,
-    coreai,
+    compression_types as _compression_types,
+    coreai as _coreai,
 )
 from coreai_opt.coreai_utils._utils.graph_utils import (
     _apply_compression_transform,
@@ -61,7 +61,7 @@
 
 
 def palettize_weights(
-    coreai_program: AIProgram,
+    coreai_program: _AIProgram,
     lut_dtype: DType | None,
     n_bits: int = 4,
     granularity: CompressionGranularity = CompressionGranularity.PER_TENSOR,
@@ -73,7 +73,7 @@ def palettize_weights(
     enable_fast_kmeans_mode: bool = True,
     rounding_precision: int = 4,
     in_place: bool = False,
-) -> AIProgram:
+) -> _AIProgram:
     """Palettize weights in a Core AI AIProgram (MLIR<CoreAI> IR) by using Core AI ops.
 
     Walks through the IR and palettizes each coreai.constant op that needs to be
@@ -178,7 +178,7 @@ def replace_weight_with_compression_op(
             fused zero_point = lut_zero_point * per_channel_scale.
         """
         if not _should_compress_op(op, weight_num_threshold, _OPS_WEIGHT_NEED_COMPRESSION):
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         const_weight: Any = op
         weight = _get_constant_value_as_np_array(const_weight)
@@ -194,7 +194,7 @@ def replace_weight_with_compression_op(
                 "The `cluster_dim` is invalid for %s. Skipped this op.",
                 const_weight.name,
             )
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         if enable_per_channel_scale:
             # Normalize by per channel scales before doing palettization.
@@ -222,7 +222,7 @@ def replace_weight_with_compression_op(
                     "Cannot perform palettization on %s. Skipped this op.",
                     const_weight.name,
                 )
-                return WalkResult.ADVANCE
+                return _WalkResult.ADVANCE
         except ImportError:
             raise
         except Exception as e:
@@ -232,12 +232,12 @@ def replace_weight_with_compression_op(
                 const_weight.name,
                 e,
             )
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
-        with const_weight.context, const_weight.location, InsertionPoint(const_weight):
+        with const_weight.context, const_weight.location, _InsertionPoint(const_weight):
             indices = _create_constant_value_from_np_array(
                 lut_params.indices,  # same shape as weight tensor (for many cases)
-                IntegerType.get_unsigned(n_bits),
+                _IntegerType.get_unsigned(n_bits),
             )
 
             vector_axis = _create_constant_value_from_np_array(
@@ -246,11 +246,11 @@ def replace_weight_with_compression_op(
                     if lut_params.vector_axis is not None
                     else np.int16(0)
                 ),
-                IntegerType.get_signed(16),
+                _IntegerType.get_signed(16),
             )
 
             weight_float_mlir_type = (
-                F32Type.get() if original_weight_type == np.float32 else F16Type.get()
+                _F32Type.get() if original_weight_type == np.float32 else _F16Type.get()
             )
 
             if lut_dtype is not None:
@@ -259,12 +259,12 @@ def replace_weight_with_compression_op(
                 weight_element_type = cast("Any", const_weight.result.type).element_type
 
                 if lut_dtype.is_int():
-                    lut_dtype_builtin = compression_types.string_to_builtin(lut_dtype)
+                    lut_dtype_builtin = _compression_types.string_to_builtin(lut_dtype)
                     ref_mlir_type = _get_string_to_mlir_type()[lut_dtype]
                     quantized_mlir_type = (
-                        IntegerType.get_signed(ref_mlir_type.width)
+                        _IntegerType.get_signed(ref_mlir_type.width)
                         if ref_mlir_type.is_signed
-                        else IntegerType.get_unsigned(ref_mlir_type.width)
+                        else _IntegerType.get_unsigned(ref_mlir_type.width)
                     )
 
                     quant_params = _compute_qparams_by_dtype(
@@ -278,7 +278,7 @@ def replace_weight_with_compression_op(
                             "Failed to compute quantization parameters for %s. Skipped this op.",
                             const_weight.name,
                         )
-                        return WalkResult.ADVANCE
+                        return _WalkResult.ADVANCE
 
                     quantized_lut_data, lut_scale, lut_zero_point = quant_params
                     if lut_zero_point is None:
@@ -335,7 +335,7 @@ def replace_weight_with_compression_op(
                             "Failed to compute quantization parameters for %s. Skipped this op.",
                             const_weight.name,
                         )
-                        return WalkResult.ADVANCE
+                        return _WalkResult.ADVANCE
 
                     quantized_lut_data, lut_scale, _ = quant_params
 
@@ -348,15 +348,15 @@ def replace_weight_with_compression_op(
                         dtype=lut_scale.dtype,
                     )
 
-                    tensor_type = RankedTensorType.get(
+                    tensor_type = _RankedTensorType.get(
                         list(quantized_lut_data.shape), fp8_mlir_type
                     )
-                    lut_quantized_attr = DenseResourceElementsAttr.get_from_buffer(
+                    lut_quantized_attr = _DenseResourceElementsAttr.get_from_buffer(
                         quantized_lut_data,
                         "dense_resource",
                         tensor_type,
                     )
-                    lut_quantized = cast("Any", coreai.ConstantOp(value=lut_quantized_attr).result)
+                    lut_quantized = cast("Any", _coreai.ConstantOp(value=lut_quantized_attr).result)
 
                     lut_scale_const = _create_constant_value_from_np_array(
                         lut_scale_reshaped,
@@ -368,23 +368,23 @@ def replace_weight_with_compression_op(
                         weight_element_type,
                     )
 
-                compressed_weight_quantized = coreai.lut_to_dense(
+                compressed_weight_quantized = _coreai.lut_to_dense(
                     indices=indices,
                     lut=lut_quantized,
                     axis=vector_axis,
                 )
                 # Cast lut_to_dense output from quantized type to float so all
                 # blockwise_shift_scale operands share the same element type.
-                cast_type = RankedTensorType.get(
+                cast_type = _RankedTensorType.get(
                     cast("Any", compressed_weight_quantized.type).shape,
                     weight_float_mlir_type,
                 )
                 compressed_weight_float = cast(
                     "Any",
-                    coreai.CastOp(cast_type, compressed_weight_quantized).result,
+                    _coreai.CastOp(cast_type, compressed_weight_quantized).result,
                 )
                 lut_t = cast("Any", lut_scale_const.type)
-                compressed_weight = coreai.blockwise_shift_scale(
+                compressed_weight = _coreai.blockwise_shift_scale(
                     data=compressed_weight_float,
                     scale=lut_scale_const,
                     offset1=lut_zero_point_const,
@@ -397,7 +397,7 @@ def replace_weight_with_compression_op(
                 # (e.g. coreai.transpose) see the same type contract.
                 compressed_weight = cast(
                     "Any",
-                    coreai.CastOp(
+                    _coreai.CastOp(
                         cast("Any", const_weight.result.type),
                         compressed_weight,
                     ).result,
@@ -409,7 +409,7 @@ def replace_weight_with_compression_op(
                     cast("Any", const_weight.result.type).element_type,
                 )
 
-                compressed_weight = coreai.lut_to_dense(
+                compressed_weight = _coreai.lut_to_dense(
                     indices=indices,
                     lut=lut,
                     axis=vector_axis,
@@ -426,15 +426,15 @@ def replace_weight_with_compression_op(
                     )
                     compressed_weight_float = cast(
                         "Any",
-                        coreai.CastOp(
-                            RankedTensorType.get(
+                        _coreai.CastOp(
+                            _RankedTensorType.get(
                                 cast("Any", compressed_weight.type).shape,
                                 weight_float_mlir_type,
                             ),
                             compressed_weight,
                         ).result,
                     )
-                    compressed_weight = coreai.blockwise_shift_scale(
+                    compressed_weight = _coreai.blockwise_shift_scale(
                         data=compressed_weight_float,
                         scale=scale,
                         offset1=zero_point,
@@ -445,15 +445,15 @@ def replace_weight_with_compression_op(
                     )
                     compressed_weight = cast(
                         "Any",
-                        coreai.CastOp(
+                        _coreai.CastOp(
                             cast("Any", const_weight.result.type),
                             compressed_weight,
                         ).result,
                     )
 
         const_weight.result.replace_all_uses_with(compressed_weight)
 
-        return WalkResult.ADVANCE
+        return _WalkResult.ADVANCE
 
     return _apply_compression_transform(
         coreai_program,