ai-edge-quantizer-nightly 0.0.1.dev20250115__py3-none-any.whl

ai_edge_quantizer/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Init file for the AI Edge quantizer package."""
+
+# pylint: disable=unused-import
+from .quantizer import Quantizer

ai_edge_quantizer/algorithm_manager.py

@@ -0,0 +1,167 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Quantizer Algorithm Manager Interface."""
+
+import enum
+from ai_edge_quantizer import algorithm_manager_api
+from ai_edge_quantizer import default_policy
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.nonlinear_quantize import float_casting
+from ai_edge_quantizer.algorithms.uniform_quantize import naive_min_max_quantize
+
+_TFLOpName = qtyping.TFLOperationName
+
+_alg_manager_instance = algorithm_manager_api.AlgorithmManagerApi()
+
+# Expose instance functions.
+get_quantization_func = _alg_manager_instance.get_quantization_func
+get_supported_ops = _alg_manager_instance.get_supported_ops
+get_init_qsv_func = _alg_manager_instance.get_init_qsv_func
+register_op_quant_config_validation_func = (
+    _alg_manager_instance.register_op_quant_config_validation_func
+)
+register_config_check_policy_func = (
+    _alg_manager_instance.register_config_check_policy
+)
+register_quantized_op = _alg_manager_instance.register_quantized_op
+is_op_registered = _alg_manager_instance.is_op_registered
+is_algorithm_registered = _alg_manager_instance.is_algorithm_registered
+check_op_quantization_config = (
+    _alg_manager_instance.check_op_quantization_config
+)
+
+
+# Quantization algorithms.
+class AlgorithmName(str, enum.Enum):
+  NO_QUANTIZE = "no_quantize"
+  MIN_MAX_UNIFORM_QUANT = naive_min_max_quantize.ALGORITHM_KEY
+  FLOAT_CASTING = float_casting.ALGORITHM_KEY
+
+
+# Register MIN_MAX_UNIFORM_QUANT algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.MIN_MAX_UNIFORM_QUANT,
+    naive_min_max_quantize.check_op_quantization_config,
+)
+
+# Register a config check policy for MIN_MAX_UNIFORM_QUANT algorithm.
+register_config_check_policy_func(
+    AlgorithmName.MIN_MAX_UNIFORM_QUANT,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+
+for op_name, materialize_func in zip(
+    (
+        _TFLOpName.INPUT,
+        _TFLOpName.OUTPUT,
+        _TFLOpName.FULLY_CONNECTED,
+        _TFLOpName.BATCH_MATMUL,
+        _TFLOpName.CONV_2D,
+        _TFLOpName.DEPTHWISE_CONV_2D,
+        _TFLOpName.CONV_2D_TRANSPOSE,
+        _TFLOpName.RESHAPE,
+        _TFLOpName.AVERAGE_POOL_2D,
+        _TFLOpName.EMBEDDING_LOOKUP,
+        _TFLOpName.SOFTMAX,
+        _TFLOpName.TANH,
+        _TFLOpName.TRANSPOSE,
+        _TFLOpName.GELU,
+        _TFLOpName.ADD,
+        _TFLOpName.SUB,
+        _TFLOpName.MUL,
+        _TFLOpName.MEAN,
+        _TFLOpName.RSQRT,
+        _TFLOpName.CONCATENATION,
+        _TFLOpName.STRIDED_SLICE,
+        _TFLOpName.SPLIT,
+        _TFLOpName.LOGISTIC,  # Sigmoid
+        _TFLOpName.SLICE,
+        _TFLOpName.SUM,
+        _TFLOpName.SELECT_V2,
+    ),
+    (
+        naive_min_max_quantize.materialize_input,
+        naive_min_max_quantize.materialize_output,
+        naive_min_max_quantize.materialize_fc_conv,
+        naive_min_max_quantize.materialize_batch_matmul,
+        naive_min_max_quantize.materialize_fc_conv,
+        naive_min_max_quantize.materialize_fc_conv,
+        naive_min_max_quantize.materialize_conv2d_transpose,
+        naive_min_max_quantize.materialize_reshape,
+        naive_min_max_quantize.materialize_average_pool_2d,
+        naive_min_max_quantize.materialize_embedding_lookup,
+        naive_min_max_quantize.materialize_softmax_and_logistic,
+        naive_min_max_quantize.materialize_tanh,
+        naive_min_max_quantize.materialize_transpose,
+        naive_min_max_quantize.materialize_gelu,
+        naive_min_max_quantize.materialize_add,
+        naive_min_max_quantize.materialize_sub,
+        naive_min_max_quantize.materialize_mul,
+        naive_min_max_quantize.materialize_mean,
+        naive_min_max_quantize.materialize_rsqrt,
+        naive_min_max_quantize.materialize_concatenation,
+        naive_min_max_quantize.materialize_strided_slice,
+        naive_min_max_quantize.materialize_split,
+        naive_min_max_quantize.materialize_softmax_and_logistic,
+        naive_min_max_quantize.materialize_slice,
+        naive_min_max_quantize.materialize_sum,
+        naive_min_max_quantize.materialize_select_v2,
+    ),
+):
+  register_quantized_op(
+      AlgorithmName.MIN_MAX_UNIFORM_QUANT,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=materialize_func,
+  )
+
+# Register FLOAT_CASTING algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.FLOAT_CASTING,
+    float_casting.check_op_quantization_config,
+)
+
+# Register a config check policy for FLOAT_CASTING algorithm.
+# TODO: b/353780772 - Replace an empty policy for FLOAT_CASTING algorithm.
+register_config_check_policy_func(
+    AlgorithmName.FLOAT_CASTING, qtyping.ConfigCheckPolicyDict()
+)
+
+for op_name, materialize_func in zip(
+    (
+        _TFLOpName.FULLY_CONNECTED,
+        _TFLOpName.CONV_2D,
+        _TFLOpName.DEPTHWISE_CONV_2D,
+        _TFLOpName.CONV_2D_TRANSPOSE,
+        _TFLOpName.EMBEDDING_LOOKUP,
+    ),
+    (
+        float_casting.materialize_fc_conv,
+        float_casting.materialize_fc_conv,
+        float_casting.materialize_fc_conv,
+        float_casting.materialize_conv2d_transpose,
+        float_casting.materialize_embedding_lookup,
+    ),
+):
+  register_quantized_op(
+      AlgorithmName.FLOAT_CASTING,
+      op_name,
+      float_casting.init_qsvs,
+      calibration_func=float_casting.calibrate,
+      materialize_func=materialize_func,
+  )

ai_edge_quantizer/algorithm_manager_api.py

@@ -0,0 +1,271 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""The Python API for Algorithm Manager of Quantizer."""
+
+from collections.abc import Callable
+import dataclasses
+import functools
+from typing import Any, Optional
+from ai_edge_quantizer import qtyping
+
+
+@dataclasses.dataclass
+class QuantizedOperationInfo:
+  """Stores all quantization functions for a given op."""
+
+  tfl_op_key: qtyping.TFLOperationName
+  init_qsv_func: Callable[..., Any]
+  calibration_func: Callable[..., Any]
+  materialize_func: Callable[..., Any]
+
+
+@dataclasses.dataclass
+class QuantizationAlgorithmInfo:
+  quantization_algorithm: str
+  quantized_ops: dict[qtyping.TFLOperationName, QuantizedOperationInfo]
+
+
+class AlgorithmManagerApi:
+  """Quantizer API client to manage quantization configs and functions."""
+
+  def __init__(self):
+    self._algorithm_registry = dict()
+    # Check if an op quantization config is supported for a given algorithm.
+    self._config_check_registry = dict()
+    # Policy to check if an op quantization config is supported for a given
+    # algorithm.
+    self._config_check_policy_registry = dict()
+
+  def register_op_quant_config_validation_func(
+      self,
+      algorithm_key: str,
+      config_check_func: Callable[..., Any],
+  ):
+    """Register functions to check if an op quantization config is supported."""
+    self._config_check_registry[algorithm_key] = config_check_func
+
+  def register_quantized_op(
+      self,
+      algorithm_key: str,
+      tfl_op_name: qtyping.TFLOperationName,
+      init_qsv_func: Callable[..., Any],
+      calibration_func: Callable[..., Any],
+      materialize_func: Callable[..., Any],
+  ):
+    """Register functions to support a quantization operation.
+
+    This function registers the relevant information to support the quantized
+    version of given tfl_operation, for the algorithm specified by
+    quantization_algorithm.
+
+    Args:
+      algorithm_key: Quantization algorithm keyword for which the quantized
+        operation is for.
+      tfl_op_name: TFLite op name.
+      init_qsv_func: QSV init function to be called.
+      calibration_func: Quantized operation to be called during calibration.
+      materialize_func: Quantized operation to be called during materialization.
+    """
+    quantized_algorithm_info = self._algorithm_registry.setdefault(
+        algorithm_key, QuantizationAlgorithmInfo(algorithm_key, dict())
+    )
+
+    quantized_algorithm_info.quantized_ops[tfl_op_name] = (
+        QuantizedOperationInfo(
+            tfl_op_name,
+            init_qsv_func,
+            calibration_func,
+            materialize_func,
+        )
+    )
+
+  def is_op_registered(
+      self,
+      quantization_algorithm: str,
+      tfl_op_name: qtyping.TFLOperationName,
+  ) -> bool:
+    """Check if the given key for quantization is valid.
+
+    Args:
+      quantization_algorithm: Target quantization algorithm.
+      tfl_op_name: TFL operation name.
+
+    Returns:
+      True if the given op is registered for the given algorithm, false
+      otherwise.
+    """
+    if not self.is_algorithm_registered(quantization_algorithm):
+      return False
+
+    return (
+        tfl_op_name
+        in self._algorithm_registry[quantization_algorithm].quantized_ops
+    )
+
+  def is_algorithm_registered(self, quantization_algorithm: str) -> bool:
+    """Check if the given algorithm is registered.
+
+    Args:
+      quantization_algorithm: Target quantization algorithm.
+
+    Returns:
+      True if the given algorithm is registered, false otherwise.
+    """
+    return quantization_algorithm in self._algorithm_registry
+
+  def check_op_quantization_config(
+      self,
+      quantization_algorithm: str,
+      tfl_op_name: qtyping.TFLOperationName,
+      op_quantization_config: qtyping.OpQuantizationConfig,
+  ) -> None:
+    """Checks if the given op quantization config is valid.
+
+    Args:
+      quantization_algorithm: Target quantization algorithm.
+      tfl_op_name: TFL operation name.
+      op_quantization_config: Op quantization config to be checked.
+
+    Raises:
+      ValueError if the given op is not registered for the given algorithm, or
+      the given algorithm is not registered.
+    """
+    if op_quantization_config.skip_checks:
+      return
+    if not self.is_op_registered(quantization_algorithm, tfl_op_name):
+      raise ValueError(
+          f"Unsupported operation {tfl_op_name} for Algorithm:"
+          f" {quantization_algorithm}."
+      )
+    if quantization_algorithm not in self._config_check_registry:
+      raise ValueError(
+          f"Config checking function for  algorithm {quantization_algorithm} is"
+          " not registered. Please use"
+          " `register_op_quant_config_validation_func` to register the"
+          " validation function."
+      )
+    self._config_check_registry[quantization_algorithm](
+        tfl_op_name,
+        op_quantization_config,
+        self._config_check_policy_registry[quantization_algorithm],
+    )
+
+  def get_supported_ops(self, alg_key: str) -> list[qtyping.TFLOperationName]:
+    """Returns the list of supported ops for the given algorithm.
+
+    Args:
+      alg_key: Algorithm key.
+
+    Returns:
+      The list of supported JAX operations.
+
+    Raises:
+      ValueError if the alg_key is not registered.
+    """
+    if alg_key not in self._algorithm_registry:
+      raise ValueError(f"Unregistered algorithm: {alg_key}")
+
+    return list(self._algorithm_registry[alg_key].quantized_ops.keys())
+
+  def get_quantization_func(
+      self,
+      algorithm_key: str,
+      tfl_op_name: qtyping.TFLOperationName,
+      quantize_mode: qtyping.QuantizeMode,
+  ) -> Callable[..., Any]:
+    """Gets the quantization function.
+
+    Args:
+      algorithm_key: Target quantization algorithm key (e.g.,
+        AlgorithmName.MIN_MAX_UNIFORM_QUANT).
+      tfl_op_name: TFLite op name.
+      quantize_mode: Quantization mode to be used.
+
+    Returns:
+      A quantized operation (function) corresponds to the requested algorithm
+      for the TFL op.
+    """
+    if not self.is_op_registered(algorithm_key, tfl_op_name):
+      raise ValueError(
+          f"Unsupported operation {tfl_op_name} for Algorithm: {algorithm_key}."
+          f" Supported ops for algorithm {algorithm_key}:"
+          f" {self.get_supported_ops(algorithm_key)}"
+      )
+
+    quantized_algorithm_info = self._algorithm_registry[algorithm_key]
+    quantized_op_info = quantized_algorithm_info.quantized_ops
+    quantized_func = self._get_target_func(
+        quantized_op_info, tfl_op_name, quantize_mode
+    )
+    if quantized_func is None:
+      raise ValueError(
+          "Cannot retrieve appropriate quantization function for"
+          f" {tfl_op_name} for algorithm {algorithm_key} under quantization"
+          f" mode {quantize_mode}. Check if the op is registed in"
+          " algorithm_manager."
+      )
+
+    return quantized_func
+
+  def get_init_qsv_func(
+      self,
+      algorithm_key: str,
+      tfl_op_name: qtyping.TFLOperationName,
+  ) -> functools.partial:
+    """Gets the initial Quantization Statistics Variable function for a given op.
+
+    Args:
+      algorithm_key: Quantization algorithm to search.
+      tfl_op_name: Target TFL operation.
+
+    Returns:
+      A function for qsv initialization.
+    """
+
+    if not self.is_op_registered(algorithm_key, tfl_op_name):
+      raise ValueError(
+          f"Unsupported operation {tfl_op_name} for Algorithm: {algorithm_key}."
+          f" Supported ops for algorithm {algorithm_key}:"
+          f" {self.get_supported_ops(algorithm_key)}"
+      )
+    quantized_algorithm_info = self._algorithm_registry[algorithm_key]
+    quantized_op_info = quantized_algorithm_info.quantized_ops
+
+    return quantized_op_info[tfl_op_name].init_qsv_func
+
+  def _get_target_func(
+      self,
+      quantized_op_info,
+      tfl_op_name: qtyping.TFLOperationName,
+      quantize_mode: qtyping.QuantizeMode,
+  ):
+    """Gets the function corresponding to the given JAX quantization phase and op."""
+    if quantize_mode == qtyping.QuantizeMode.CALIBRATE:
+      return quantized_op_info[tfl_op_name].calibration_func
+    elif quantize_mode == qtyping.QuantizeMode.MATERIALIZE:
+      return quantized_op_info[tfl_op_name].materialize_func
+    return None
+
+  # TODO: b/53780772 - Merge this function with
+  # register_op_quant_config_validation_func after full transition to new check
+  # mechanism.
+  def register_config_check_policy(
+      self,
+      algorithm_key: str,
+      config_check_policy: Optional[qtyping.ConfigCheckPolicyDict],
+  ):
+    """Registers a policy to check the op quantization config."""
+    self._config_check_policy_registry[algorithm_key] = config_check_policy

ai_edge_quantizer/algorithm_manager_api_test.py

@@ -0,0 +1,210 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for algorithm_manager_api."""
+
+from absl.testing import parameterized
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import algorithm_manager_api
+from ai_edge_quantizer import qtyping
+
+_TFLOpName = qtyping.TFLOperationName
+
+
+# Sample functions for test cases.
+def _sample_init_qsvs(*_, **__):
+  return 1.0, dict()
+
+
+def _sample_calibration_func(*_, **__):
+  return 2.0, dict()
+
+
+def _sample_materialize_func(*_, **__):
+  return 3.0, dict()
+
+
+def _sample_check_op_config_func(_, op_config):
+  if op_config.weight_tensor_config.num_bits == 17:
+    raise ValueError("Unsupported number of bits.")
+
+
+class AlgorithmManagerApiTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._alg_manager = algorithm_manager_api.AlgorithmManagerApi()
+
+  def test_register_op_quant_config_validation_func_succeeds(self):
+    self.assertEmpty(self._alg_manager._config_check_registry)
+    test_algorithm_name = "test_algorithm"
+    self._alg_manager.register_op_quant_config_validation_func(
+        test_algorithm_name, _sample_check_op_config_func
+    )
+    self.assertIn(test_algorithm_name, self._alg_manager._config_check_registry)
+    check_func = self._alg_manager._config_check_registry[test_algorithm_name]
+    self.assertEqual(check_func, _sample_check_op_config_func)
+
+  def test_register_quantized_op(self):
+    self._alg_manager.register_quantized_op(
+        algorithm_key="ptq",
+        tfl_op_name=_TFLOpName.FULLY_CONNECTED,
+        init_qsv_func=_sample_init_qsvs,
+        calibration_func=_sample_calibration_func,
+        materialize_func=_sample_materialize_func,
+    )
+    self._alg_manager.register_quantized_op(
+        algorithm_key="gptq",
+        tfl_op_name=_TFLOpName.CONV_2D,
+        init_qsv_func=_sample_init_qsvs,
+        calibration_func=_sample_calibration_func,
+        materialize_func=_sample_materialize_func,
+    )
+    self.assertTrue(self._alg_manager.is_algorithm_registered("ptq"))
+    self.assertTrue(self._alg_manager.is_algorithm_registered("gptq"))
+    self.assertTrue(
+        self._alg_manager.is_op_registered("ptq", _TFLOpName.FULLY_CONNECTED)
+    )
+    self.assertTrue(
+        self._alg_manager.is_op_registered("gptq", _TFLOpName.CONV_2D)
+    )
+    self.assertFalse(
+        self._alg_manager.is_op_registered("gptq", _TFLOpName.DEPTHWISE_CONV_2D)
+    )
+
+  def test_get_supported_ops(self):
+    algorithm_key = "ptq"
+    self._alg_manager.register_quantized_op(
+        algorithm_key=algorithm_key,
+        tfl_op_name=_TFLOpName.FULLY_CONNECTED,
+        init_qsv_func=_sample_init_qsvs,
+        calibration_func=_sample_calibration_func,
+        materialize_func=_sample_materialize_func,
+    )
+    self._alg_manager.register_quantized_op(
+        algorithm_key=algorithm_key,
+        tfl_op_name=_TFLOpName.CONV_2D,
+        init_qsv_func=_sample_init_qsvs,
+        calibration_func=_sample_calibration_func,
+        materialize_func=_sample_materialize_func,
+    )
+    supported_ops = self._alg_manager.get_supported_ops(algorithm_key)
+    self.assertIn(_TFLOpName.CONV_2D, supported_ops)
+    self.assertIn(_TFLOpName.FULLY_CONNECTED, supported_ops)
+    self.assertNotIn(_TFLOpName.DEPTHWISE_CONV_2D, supported_ops)
+
+  def test_get_quantization_func(self):
+    algorithm_key = "ptq"
+    tfl_op = _TFLOpName.FULLY_CONNECTED
+    self._alg_manager.register_quantized_op(
+        algorithm_key=algorithm_key,
+        tfl_op_name=tfl_op,
+        init_qsv_func=_sample_init_qsvs,
+        calibration_func=_sample_calibration_func,
+        materialize_func=_sample_materialize_func,
+    )
+    materialize_func = self._alg_manager.get_quantization_func(
+        algorithm_key,
+        tfl_op,
+        qtyping.QuantizeMode.MATERIALIZE,
+    )
+    self.assertEqual(_sample_materialize_func()[0], materialize_func()[0])
+    calibration_func = self._alg_manager.get_quantization_func(
+        algorithm_key,
+        tfl_op,
+        qtyping.QuantizeMode.CALIBRATE,
+    )
+    self.assertEqual(_sample_calibration_func()[0], calibration_func()[0])
+
+    # Query for unsupported operation.
+    error_message = "Unsupported operation"
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      self._alg_manager.get_quantization_func(
+          algorithm_key,
+          _TFLOpName.BATCH_MATMUL,
+          qtyping.QuantizeMode.MATERIALIZE,
+      )
+
+    # Query for unregisted algorithm.
+    error_message = "Unregistered algorithm"
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      self._alg_manager.get_quantization_func(
+          "gptq",
+          tfl_op,
+          qtyping.QuantizeMode.MATERIALIZE,
+      )
+
+  def test_get_init_qsv_func(self):
+    algorithm_key = "ptq"
+    tfl_op = _TFLOpName.FULLY_CONNECTED
+    self._alg_manager.register_quantized_op(
+        algorithm_key=algorithm_key,
+        tfl_op_name=tfl_op,
+        init_qsv_func=_sample_init_qsvs,
+        calibration_func=_sample_calibration_func,
+        materialize_func=_sample_materialize_func,
+    )
+    init_qsv_func = self._alg_manager.get_init_qsv_func(algorithm_key, tfl_op)
+    self.assertEqual(_sample_init_qsvs()[0], init_qsv_func()[0])
+
+    # Query for unsupported operation.
+    error_message = "Unsupported operation"
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      self._alg_manager.get_init_qsv_func(
+          algorithm_key,
+          _TFLOpName.BATCH_MATMUL,
+      )
+
+    # Query for unregisted algorithm.
+    error_message = "Unregistered algorithm"
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      self._alg_manager.get_init_qsv_func(
+          "gptq",
+          tfl_op,
+      )
+
+  def test_register_config_check_policy_succeeds(self):
+    self.assertEmpty(self._alg_manager._config_check_policy_registry)
+    test_algorithm_name = "test_algorithm"
+    test_config_check_policy = qtyping.ConfigCheckPolicyDict({
+        _TFLOpName.FULLY_CONNECTED: {
+            qtyping.OpQuantizationConfig(
+                weight_tensor_config=qtyping.TensorQuantizationConfig(
+                    num_bits=1
+                )
+            )
+        }
+    })
+    self._alg_manager.register_config_check_policy(
+        test_algorithm_name, test_config_check_policy
+    )
+    self.assertIn(
+        test_algorithm_name, self._alg_manager._config_check_policy_registry
+    )
+    self.assertIsNotNone(
+        self._alg_manager._config_check_policy_registry[test_algorithm_name]
+    )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/algorithms/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+

ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+