improve einsum

Author: JyotinderSingh

keras/src/layers/core/dense.py

@@ -9,6 +9,7 @@
 from keras.src import quantizers
 from keras.src import regularizers
 from keras.src.api_export import keras_export
+from keras.src.dtype_policies import QuantizedFloat8DTypePolicy
 from keras.src.layers.input_spec import InputSpec
 from keras.src.layers.layer import Layer
 from keras.src.quantizers.quantization_config import QuantizationConfig
@@ -384,8 +385,6 @@ def quantized_build(self, kernel_shape, mode, config=None):
         self._is_quantized = True
 
     def _int8_build(self, kernel_shape, config=None):
-        from keras.src.quantizers.quantization_config import QuantizationConfig
-
         # Per-channel int8 quantizer for the last axis (features).
         self.inputs_quantizer = (
             QuantizationConfig.activation_quantizer_or_default(
@@ -500,8 +499,6 @@ def _int4_build(self, kernel_shape, config=None):
         int8 byte.
         """
         # Per-channel int8 quantizer for the last axis (features).
-        from keras.src.quantizers.quantization_config import QuantizationConfig
-
         self.inputs_quantizer = (
             QuantizationConfig.activation_quantizer_or_default(
                 config, quantizers.AbsMaxQuantizer(axis=-1)
@@ -529,8 +526,6 @@ def _int4_build(self, kernel_shape, config=None):
         self._orig_input_dim = input_dim
 
     def _float8_build(self):
-        from keras.src.dtype_policies import QuantizedFloat8DTypePolicy
-
         # If `self.dtype_policy` is not QuantizedFloat8DTypePolicy, then set
         # `amax_history_length` to its default value.
         amax_history_length = getattr(
@@ -781,16 +776,6 @@ def quantize(self, mode=None, type_check=True, config=None):
 
         kernel_shape = self._kernel.shape
         if mode == "int8":
-            # Handle activation quantization
-            if config.activation_quantizer:
-                self.inputs_quantizer = config.activation_quantizer
-            elif config.activation_quantizer is None:
-                # Weight-only quantization
-                pass
-            else:
-                # Default behavior
-                self.inputs_quantizer = quantizers.AbsMaxQuantizer(axis=-1)
-
             # Handle weight quantization
             # Quantize `self._kernel` to int8 and compute corresponding scale
             weight_quantizer = QuantizationConfig.weight_quantizer_or_default(

keras/src/layers/core/einsum_dense.py

@@ -13,10 +13,10 @@
 from keras.src import quantizers
 from keras.src import regularizers
 from keras.src.api_export import keras_export
+from keras.src.dtype_policies import QuantizedFloat8DTypePolicy
 from keras.src.layers.input_spec import InputSpec
 from keras.src.layers.layer import Layer
 from keras.src.quantizers.quantization_config import QuantizationConfig
-from keras.src.quantizers.quantization_config import validate_and_resolve_config
 from keras.src.quantizers.quantizers import dequantize_with_sz_map
 
 
@@ -457,91 +457,8 @@ def quantized_build(self, kernel_shape, mode, config=None):
             raise self._quantization_mode_error(mode)
         self._is_quantized = True
 
-    def quantize(self, mode=None, type_check=True, config=None):
-        # Prevent quantization of the subclasses
-        if type_check and (type(self) is not EinsumDense):
-            raise self._not_implemented_error(self.quantize)
-
-        config = validate_and_resolve_config(mode, config)
-        mode = config.mode
-
-        kernel_shape = self._kernel.shape
-        if mode == "int8":
-            # Handle activation quantization
-            if config.activation_quantizer:
-                self.inputs_quantizer = config.activation_quantizer
-            elif config.activation_quantizer is None:
-                # Weight-only quantization
-                pass
-            else:
-                # Default behavior
-                self.inputs_quantizer = quantizers.AbsMaxQuantizer(
-                    axis=self._input_reduced_axes
-                )
-
-            # Handle weight quantization
-            # Quantize `self._kernel` to int8 and compute corresponding scale
-            weight_quantizer = QuantizationConfig.weight_quantizer_or_default(
-                config, quantizers.AbsMaxQuantizer(axis=0)
-            )
-            self._kernel, self.kernel_scale = weight_quantizer(
-                self._kernel, to_numpy=True
-            )
-            self.quantized_build(kernel_shape, mode, config)
-
-        elif mode == "int4":
-            # Handle activation quantization
-            if config.activation_quantizer:
-                self.inputs_quantizer = config.activation_quantizer
-            elif config.activation_quantizer is None:
-                # Weight-only quantization
-                pass
-            else:
-                # Default behavior
-                self.inputs_quantizer = quantizers.AbsMaxQuantizer(
-                    axis=self._input_reduced_axes
-                )
-
-            # Handle weight quantization
-            # 1. Quantize to int4 values (stored in int8 dtype, range [-8, 7])
-            weight_quantizer = QuantizationConfig.weight_quantizer_or_default(
-                config,
-                quantizers.AbsMaxQuantizer(
-                    axis=0,
-                    value_range=(-8, 7),
-                    output_dtype="int8",
-                ),
-            )
-            self._kernel, self.kernel_scale = weight_quantizer(
-                self._kernel, to_numpy=True
-            )
-            # 2. Pack two int4 values into a single int8 byte.
-            # Choose the axis to perform int4 packing - use the first reduced
-            # axis for the kernel (analogous to the input dimension of a Dense
-            # layer).
-            self._int4_pack_axis = (
-                self._kernel_reduced_axes[0] if self._kernel_reduced_axes else 0
-            )
-            self._kernel, _, _ = quantizers.pack_int4(
-                self._kernel, axis=self._int4_pack_axis
-            )
-            self.quantized_build(kernel_shape, mode, config)
-
-        elif mode == "float8":
-            self.quantized_build(kernel_shape, mode)
-
-        elif mode == "gptq":
-            self.quantized_build(kernel_shape, mode, config)
-
-        # Set new dtype policy.
-        if self.dtype_policy.quantization_mode is None:
-            policy = dtype_policies.get(f"{mode}_from_{self.dtype_policy.name}")
-            self.dtype_policy = policy
-
     def _int8_build(self, kernel_shape, config=None):
         self._set_quantization_info()
-        from keras.src.quantizers.quantization_config import QuantizationConfig
-
         self.inputs_quantizer = (
             QuantizationConfig.activation_quantizer_or_default(
                 config,
@@ -691,8 +608,6 @@ def _int4_build(self, kernel_shape, config=None):
         self._set_quantization_info()
 
         # Quantizer for the inputs (per the reduced axes)
-        from keras.src.quantizers.quantization_config import QuantizationConfig
-
         self.inputs_quantizer = (
             QuantizationConfig.activation_quantizer_or_default(
                 config,
@@ -736,8 +651,6 @@ def _int4_build(self, kernel_shape, config=None):
         )
 
     def _float8_build(self):
-        from keras.src.dtype_policies import QuantizedFloat8DTypePolicy
-
         # If `self.dtype_policy` is not QuantizedFloat8DTypePolicy, then set
         # `amax_history_length` to its default value.
         amax_history_length = getattr(
@@ -903,7 +816,8 @@ def grad_fn(*args, upstream=None):
             # Quantize inputs per `self.inputs_quantizer`.
             if self.inputs_quantizer:
                 inputs_q, inputs_scale = self.inputs_quantizer(inputs)
-                # Align `inputs_scale` axes with the output for correct broadcasting
+                # Align `inputs_scale` axes with the output
+                # for correct broadcasting
                 inputs_scale = self._adjust_scale_for_quant(
                     inputs_scale, "input"
                 )
@@ -1036,10 +950,8 @@ def quantize(self, mode, type_check=True, config=None):
             raise self._not_implemented_error(self.quantize)
 
         kernel_shape = self._kernel.shape
-        if mode in ("int8", "int4", "gptq"):
-            self._set_quantization_info()
 
-        from keras.src.quantizers.quantization_config import QuantizationConfig
+        self._set_quantization_info()
 
         if mode == "int8":
             # Quantize `self._kernel` to int8 and compute corresponding scale

keras/src/layers/core/reversible_embedding.py

@@ -187,7 +187,6 @@ def _int8_build(self, embeddings_shape, config=None):
         if embeddings_shape is None:
             embeddings_shape = (self.input_dim, self.output_dim)
         super()._int8_build(embeddings_shape=embeddings_shape)
-        from keras.src.quantizers.quantization_config import QuantizationConfig
 
         self.inputs_quantizer = (
             QuantizationConfig.activation_quantizer_or_default(
@@ -213,7 +212,6 @@ def _int4_build(self, embeddings_shape, config=None):
         if embeddings_shape is None:
             embeddings_shape = (self.input_dim, self.output_dim)
         super()._int4_build(embeddings_shape=embeddings_shape, config=config)
-        from keras.src.quantizers.quantization_config import QuantizationConfig
 
         self.inputs_quantizer = (
             QuantizationConfig.activation_quantizer_or_default(

keras/src/models/model.py

@@ -9,6 +9,7 @@
 from keras.src.layers.layer import Layer
 from keras.src.models.variable_mapping import map_saveable_variables
 from keras.src.quantizers.gptq_core import gptq_quantize
+from keras.src.quantizers.quantization_config import validate_and_resolve_config
 from keras.src.saving import saving_api
 from keras.src.trainers import trainer as base_trainer
 from keras.src.utils import summary_utils
@@ -433,9 +434,6 @@ def quantize(self, mode=None, config=None, **kwargs):
                 time.
             config: The configuration of the quantization.
         """
-        from keras.src.quantizers.quantization_config import (
-            validate_and_resolve_config,
-        )
 
         # Validate inputs.
         type_check = kwargs.pop("type_check", True)

keras/src/quantizers/gptq_test.py

@@ -617,7 +617,8 @@ def test_quantize_gptq_combinations(self, dataset, config):
             "mode": "gptq",
             "config": {"weight_bits": 4},
             "expected_exception": ValueError,
-            "error_msg": "Argument `config` must be an instance of `QuantizationConfig`",
+            "error_msg": "Argument `config` must be an instance of "
+            "`QuantizationConfig`",
         },
         {
             "testcase_name": "gptq_with_none_config",

keras/src/quantizers/quantization_config.py

@@ -54,9 +54,9 @@ def activation_quantizer_or_default(config, default):
 @keras_export("keras.quantizers.Int8QuantizationConfig")
 class Int8QuantizationConfig(QuantizationConfig):
     def __init__(self, weight_quantizer=None, activation_quantizer="default"):
-        if activation_quantizer == "default":
-            from keras.src.quantizers.quantizers import AbsMaxQuantizer
+        from keras.src.quantizers.quantizers import AbsMaxQuantizer
 
+        if activation_quantizer == "default":
             activation_quantizer = AbsMaxQuantizer(axis=-1)
         super().__init__(weight_quantizer, activation_quantizer)
         if self.weight_quantizer:
@@ -76,9 +76,9 @@ def mode(self):
 @keras_export("keras.quantizers.Int4QuantizationConfig")
 class Int4QuantizationConfig(QuantizationConfig):
     def __init__(self, weight_quantizer=None, activation_quantizer="default"):
-        if activation_quantizer == "default":
-            from keras.src.quantizers.quantizers import AbsMaxQuantizer
+        from keras.src.quantizers.quantizers import AbsMaxQuantizer
 
+        if activation_quantizer == "default":
             activation_quantizer = AbsMaxQuantizer(axis=-1)
         super().__init__(weight_quantizer, activation_quantizer)
         if self.weight_quantizer: