reworke on review comments

Author: amitsrivastava78

keras/api/__init__.py

@@ -60,7 +60,6 @@
 from keras.src.ops.function import Function as Function
 from keras.src.ops.operation import Operation as Operation
 from keras.src.optimizers.optimizer import Optimizer as Optimizer
-from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.regularizers.regularizers import Regularizer as Regularizer
 from keras.src.version import __version__ as __version__

keras/api/_tf_keras/keras/__init__.py

@@ -58,7 +58,6 @@
 from keras.src.ops.function import Function as Function
 from keras.src.ops.operation import Operation as Operation
 from keras.src.optimizers.optimizer import Optimizer as Optimizer
-from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.regularizers.regularizers import Regularizer as Regularizer
 from keras.src.version import __version__ as __version__

keras/api/_tf_keras/keras/quantizers/__init__.py

@@ -7,7 +7,6 @@
 from keras.src.quantizers import deserialize as deserialize
 from keras.src.quantizers import get as get
 from keras.src.quantizers import serialize as serialize
-from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantizers import AbsMaxQuantizer as AbsMaxQuantizer
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.quantizers.quantizers import abs_max_quantize as abs_max_quantize

keras/api/quantizers/__init__.py

@@ -7,7 +7,6 @@
 from keras.src.quantizers import deserialize as deserialize
 from keras.src.quantizers import get as get
 from keras.src.quantizers import serialize as serialize
-from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantizers import AbsMaxQuantizer as AbsMaxQuantizer
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.quantizers.quantizers import abs_max_quantize as abs_max_quantize

keras/src/models/model.py

@@ -438,7 +438,8 @@ def quantize(self, mode, config=None, **kwargs):
             if not isinstance(config, GPTQConfig):
                 raise TypeError(
                     "When using 'gptq' mode, you must pass a `config` "
-                    "argument of type `keras.quantizers.GPTQConfig`."
+                    "argument of type "
+                    "`keras.quantizers.gptq_config.GPTQConfig`."
                 )
             # The config object's own quantize method drives the process
             config.quantize(self)

keras/src/models/model_test.py

@@ -7,6 +7,7 @@
 import pytest
 from absl.testing import parameterized
 
+from keras.quantizers import GPTQConfig
 from keras.src import backend
 from keras.src import layers
 from keras.src import losses
@@ -17,7 +18,6 @@
 from keras.src.models.functional import Functional
 from keras.src.models.model import Model
 from keras.src.models.model import model_from_json
-from keras.src.quantizers.gptq_config import GPTQConfig
 
 
 def _get_model():

keras/src/quantizers/gptq_config.py

@@ -1,10 +1,8 @@
 from absl import logging
 
-from keras.src.api_export import keras_export
 from keras.src.quantizers.gptq_core import quantize_model
 
 
-@keras_export(["keras.GPTQConfig", "keras.quantizers.GPTQConfig"])
 class GPTQConfig:
     """Configuration class for the GPTQ algorithm.
 
@@ -57,7 +55,6 @@ def __init__(
         self.group_size = group_size
         self.symmetric = symmetric
         self.act_order = act_order
-        self.quantization_method = "gptq"
 
     def quantize(self, model):
         """

keras/src/quantizers/gptq_core.py

@@ -9,7 +9,7 @@
 from keras.src.layers import EinsumDense
 from keras.src.layers import Embedding
 from keras.src.quantizers.gptq import GPTQ
-from keras.src.quantizers.gptq_quant import GPTQQuant
+from keras.src.quantizers.gptq_quant import GPTQQuantization
 
 
 def get_dataloader(tokenizer, seqlen, dataset, nsamples=128):
@@ -271,7 +271,7 @@ def hook(*args, **kwargs):
                 inp_reshaped = ops.reshape(layer_inputs, (-1, num_features))
                 gptq_object.update_hessian_with_batch(inp_reshaped)
 
-                quantizer = GPTQQuant(
+                quantizer = GPTQQuantization(
                     wbits,
                     perchannel=True,
                     symmetric=symmetric,
@@ -331,5 +331,3 @@ def quantize_model(model, config):
         config.act_order,
         config.wbits,
     )
-
-    return

keras/src/quantizers/gptq_quant.py

@@ -15,8 +15,8 @@ def dequantize(x, scale, zero, maxq):
     return ops.multiply(scale, dequantized_x)
 
 
-class GPTQQuant:
-    """Initializes the GPTQQuant state.
+class GPTQQuantization:
+    """Initializes the GPTQQuantization state.
 
     Args:
         shape (int, optional): This argument is currently unused.

keras/src/quantizers/gptq_test.py

@@ -5,7 +5,7 @@
 from keras.src import ops
 from keras.src import testing
 from keras.src.quantizers.gptq import GPTQ
-from keras.src.quantizers.gptq_quant import GPTQQuant
+from keras.src.quantizers.gptq_quant import GPTQQuantization
 
 
 def _get_mock_layer(layer_type, kernel_shape, rng):
@@ -69,7 +69,7 @@ def test_full_quantization_process(self):
         original_weights = np.copy(ops.convert_to_numpy(mock_layer.kernel))
 
         gptq_instance = GPTQ(mock_layer)
-        gptq_instance.quantizer = GPTQQuant(wbits=4, symmetric=False)
+        gptq_instance.quantizer = GPTQQuantization(wbits=4, symmetric=False)
         calibration_data = rng.standard_normal(size=(128, 16)).astype("float32")
         gptq_instance.update_hessian_with_batch(calibration_data)
         gptq_instance.quantize_and_correct_block()