Reworked on review comments

Author: amitsrivastava78

keras/api/__init__.py

@@ -60,6 +60,7 @@
 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.gptqconfig 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,6 +58,7 @@
 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.gptqconfig 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,6 +7,7 @@
 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.gptqconfig 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,6 +7,7 @@
 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.gptqconfig 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/quantizers/gptq.py

@@ -1,8 +1,7 @@
 from keras.src import ops
 from keras.src.layers import Dense
 from keras.src.layers import EinsumDense
-
-from .gptqquant import quantize
+from keras.src.quantizers.gptqquant import dequantize
 
 
 class GPTQ:
@@ -201,7 +200,7 @@ def quantize_and_correct_block(
                     )
 
                 # Quantize the current weight column
-                q = quantize(
+                q = dequantize(
                     ops.expand_dims(w, 1),
                     self.quantizer.scale,
                     self.quantizer.zero,

keras/src/quantizers/gptqconfig.py

@@ -1,8 +1,10 @@
 from absl import logging
 
-from .gptqutils import quantize_model
+from keras.src.api_export import keras_export
+from keras.src.quantizers.gptqutils import quantize_model
 
 
+@keras_export(["keras.GPTQConfig", "keras.quantizers.GPTQConfig"])
 class GPTQConfig:
     """
     Configuration class for the GPTQ (Generative Pre-trained Transformer

keras/src/quantizers/gptqquant.py

@@ -1,7 +1,7 @@
 from keras.src import ops
 
 
-def quantize(x, scale, zero, maxq):
+def dequantize(x, scale, zero, maxq):
     """The core quantization function with correct broadcasting."""
     # Ensure scale is broadcastable with the input tensor x
     if scale.shape != x.shape:
@@ -12,9 +12,12 @@ def quantize(x, scale, zero, maxq):
         zero = ops.broadcast_to(zero, x.shape)
 
     scale = ops.where(ops.equal(scale, 0), 1e-8, scale)
-    q = ops.round(x / scale) + zero
+    quantized_x = ops.divide(x, scale)
+    quantized_x = ops.round(quantized_x)
+    q = ops.add(quantized_x, zero)
     q = ops.clip(q, 0, maxq)
-    return scale * (q - zero)
+    dequantized_x = ops.subtract(q, zero)
+    return ops.multiply(scale, dequantized_x)
 
 
 class GPTQQuant:

keras/src/quantizers/gptqutils.py

@@ -11,9 +11,8 @@
 from keras.src.layers import Dense
 from keras.src.layers import EinsumDense
 from keras.src.layers import Embedding
-
-from .gptq import GPTQ
-from .gptqquant import GPTQQuant
+from keras.src.quantizers.gptq import GPTQ
+from keras.src.quantizers.gptqquant import GPTQQuant
 
 
 def get_dataloader(tokenizer, seqlen, dataset, nsamples=128, seed=0):