[JAX] Add support for per-layer config custimization and composable configs by bkowalskiINTEL · Pull Request #2481 · intel/neural-compressor

bkowalskiINTEL · 2026-05-27T15:10:27Z

No description provided.

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

for more information, see https://pre-commit.ci

Copilot

Pull request overview

Adds JAX quantization support for (1) per-layer filtering via include/exclude patterns and (2) composing multiple quantization configs, while refactoring quantize flow to centralize model wrapping and improve (de)serialization behavior.

Changes:

Add include / exclude layer filters to StaticQuantConfig and DynamicQuantConfig, and apply them when generating model info.
Add ComposableConfig support in quantization config JSON serialization/deserialization and in quantize_model() mapping construction.
Refactor JAX static/dynamic algorithms to operate on per-layer config mappings and move wrapper application into quantize_model().

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
`test/jax/test_config_on_vit.py`	Adds a ViT config-filter demonstration (currently executes at import time).
`neural_compressor/jax/quantization/saving.py`	Adds ComposableConfig (de)serialization and updates deserialization preparation to handle multiple sub-configs and include/exclude.
`neural_compressor/jax/quantization/quantize.py`	Adds ComposableConfig mapping merge logic, enforces static-first algorithm application, and centralizes model wrapping.
`neural_compressor/jax/quantization/config.py`	Introduces include/exclude filtering and switches model info to use layer paths (with filtering).
`neural_compressor/jax/algorithms/static.py`	Updates static quantization to prepare only layers selected by configs_mapping and to support per-layer params.
`neural_compressor/jax/algorithms/dynamic.py`	Updates dynamic quantization to prepare only layers selected by configs_mapping and to support per-layer params.
`examples/jax/keras/gemma/quantization.py`	Updates Gemma quantization example (currently contains an early `exit()`).

+    def _matches(pattern: str) -> bool:
+        if pattern == class_name:
+            return True
+        return re.search(pattern, layer_id) is not None


+    # Build configs_mapping - handle ComposableConfig by calling sub-configs individually
+    if isinstance(quant_config, ComposableConfig):
+        configs_mapping = _build_configs_mapping_composable(model, quant_config)
+    else:
+        model_info = quant_config.get_model_info(model)
+        configs_mapping = quant_config.to_config_mapping(model_info=model_info)


anko-intel

Some comments from offline reviewe

anko-intel · 2026-06-02T11:17:55Z


    Returns:
-        keras.Model: The quantized model wrapped for inference.
+        keras.Model: The quantized model.


Previous comment seems to be more accurate

anko-intel · 2026-06-02T12:32:12Z

+        causal_lm_make_replace_generate_function(model)
+
+    # Execute algorithms - static first to ensure calibration runs on original FP32 model
+    algo_order = sorted(algos_mapping.keys(), key=lambda name: (0 if name == STATIC_QUANT else 1))


can we use priority or something like that. I am not sure, but I think I saw such list of algos

anko-intel · 2026-06-02T12:39:03Z

-    iterate_over_layers(qmodel, operations, filter_function=lambda c: c in static_quant_mapping)
+    # Phase 1: Prepare layers and add observers
+    for layer in qmodel._flatten_layers():
+        if layer.__class__ not in static_quant_mapping:


maybe we can filter if class could be quantied in earlier stage

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Co-authored-by: Bartosz Kowalski <bartosz.kowalski@intel.com> Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Fixes missing import Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 9 comments.

+    def _matches(pattern: str) -> bool:
+        if pattern == class_name:
+            return True
+        return re.search(pattern, layer_id) is not None


+            include (Optional[List[str]]): List of layer class names or path patterns to include.
+                When set, only matching layers are quantized. Supports fnmatch patterns.
+            exclude (Optional[List[str]]): List of layer class names or path patterns to exclude.
+                Matching layers are skipped. Supports fnmatch patterns.


anko-intel · 2026-07-06T08:05:44Z

+    and the layer path (regex search).
+
+    Args:
+        layer_id: Layer path or name identifier.


maybe 2 examples also here

anko-intel · 2026-07-06T08:06:30Z

+    Args:
+        layer_id: Layer path or name identifier.
+        class_name: Layer class name (e.g. "Dense").
+        include: If set, layer must match at least one entry.


what is entry here ?

anko-intel · 2026-07-06T08:23:57Z

+    """Check if a layer passes include/exclude filters.
+
+    Each filter entry is matched against both the layer class name (exact match)
+    and the layer path (regex search).


some more description? only one of include or exclude is optional ? which one is more important/ win

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

Comments suppressed due to low confidence (1)

neural_compressor/jax/quantization/saving.py:41

The return-shape description is now inaccurate for composable configs: the function returns {quantization_type: \"composable\", configs: [...]} (note configs, not config). Updating the docstring to reflect the exact keys/shape will avoid confusion for consumers parsing saved artifacts.

def quant_config_to_json_object(quant_config: BaseConfig) -> dict:
    """Serialize a quant config to a JSON-compatible dict with class name.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 9 comments.

 from typing import Callable, Dict, Tuple

 import keras
+from keras_hub.src.models.causal_lm import CausalLM



 @register_config(framework_name=FRAMEWORK_NAME, algo_name=STATIC_QUANT)
-class StaticQuantConfig(BaseConfig):
+class StaticQuantConfig(JaxBaseConfig):


    name = STATIC_QUANT

-    def __init__(
-        self,
-        weight_dtype: str = "fp8_e4m3",
-        activation_dtype: str = "fp8_e4m3",
-        const_scale: bool = False,
-        const_weight: bool = False,
-        white_list: Optional[List[OP_NAME_OR_MODULE_TYPE]] = DEFAULT_WHITE_LIST,
-    ):
-        """Init Static quantization config.
-
-        Args:
-            weight_dtype (str): Data type for weights, default is "fp8_e4m3".
-            activation_dtype (str): Data type for activations, default is "fp8_e4m3".
-            const_scale (bool): Whether to use a constant scale factor for quantization.
-            const_weight (bool): Whether to use constant quantized weights.
-            white_list (list): A list of supported operators of this algorithm.
-
-        Returns:
-            None: Initializes the configuration instance.
-        """
-        super().__init__(white_list=white_list)
-        if not isinstance(weight_dtype, list):
-            jnp_weight_dtype = dtype_mapping[weight_dtype]
-            jnp_activation_dtype = dtype_mapping[activation_dtype]
-            if (
-                jnp.issubdtype(jnp_weight_dtype, jnp.floating) and jnp.issubdtype(jnp_activation_dtype, jnp.integer)
-            ) or (jnp.issubdtype(jnp_weight_dtype, jnp.integer) and jnp.issubdtype(jnp_activation_dtype, jnp.floating)):
-                raise ValueError("Mixed quantization with floating-point and integer dtypes is not supported.")
+    def _get_supported_class_names(self) -> set:


+        def _matches(pattern) -> bool:
+            pattern = pattern if isinstance(pattern, str) else pattern.__name__
+            if pattern == class_name:
+                return True
+            try:
+                return re.search(pattern, layer_id) is not None
+            except re.error as e:
+                raise ValueError(f"Invalid regex pattern {pattern!r} in white_list/exclude filter.") from e


wpietka · 2026-07-08T11:47:50Z

        const_scale: bool = False,
        const_weight: bool = False,
        white_list: Optional[List[OP_NAME_OR_MODULE_TYPE]] = DEFAULT_WHITE_LIST,
+        exclude: Optional[List[str]] = None,


Maybe now we should name it black_list, since we use white_list instead of include

bkowalskiINTEL · 2026-07-08T13:22:14Z

+                raise ValueError(f"Invalid regex pattern {pattern!r} in white_list/exclude filter.") from e
+
+        white_list = self._white_list
+        if white_list != DEFAULT_WHITE_LIST and (white_list is None or not any(_matches(p) for p in white_list)):


Make this if statement clearer, it's hard to understand what's going on

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

bkowalskiINTEL and others added 9 commits May 27, 2026 08:08

add include/exclude filtering and path-based get_model_info

aed7004

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

refactor algorithms to use per-layer config from configs_mapping

09c989c

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

support ComposableConfig and centralize model wrapping

3b1aa05

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

handle ComposableConfig serialization and per-layer deserialization

24ba943

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

add ViT config filtering integration test

bc8d9a6

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

update gemma example for composable config API

6b1f2b1

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

simplify deserialization to use class-based layer matching

8835c1c

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Fix configs merging

306286e

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

0f17bea

for more information, see https://pre-commit.ci

anko-intel reviewed May 29, 2026

View reviewed changes

Comment thread examples/jax/keras/gemma/quantization.py Outdated

anko-intel requested a review from Copilot June 2, 2026 09:07

Copilot started reviewing on behalf of anko-intel June 2, 2026 09:07 View session

Copilot AI reviewed Jun 2, 2026

View reviewed changes

anko-intel reviewed Jun 2, 2026

View reviewed changes

bkowalskiINTEL added 2 commits June 3, 2026 06:33

Align dynamic algo to static algo and fix white list bug

c3cdd3f

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Merge branch 'master' into dev/bkowalsk/jax_composable_configs

f8b4ad3

bkowalskiINTEL commented Jun 3, 2026

View reviewed changes

Comment thread neural_compressor/jax/algorithms/dynamic.py Outdated

Remove debug stuff commited by accident

c8e0698

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

bkowalskiINTEL commented Jun 3, 2026

View reviewed changes

Comment thread neural_compressor/jax/algorithms/static.py Outdated

bkowalskiINTEL commented Jun 3, 2026

View reviewed changes

Comment thread neural_compressor/jax/algorithms/dynamic.py Outdated

bkowalskiINTEL commented Jun 3, 2026

View reviewed changes

Comment thread neural_compressor/jax/quantization/quantize.py

bkowalskiINTEL force-pushed the dev/bkowalsk/jax_composable_configs branch from 2c557c7 to c45031b Compare June 3, 2026 14:03

bkowalskiINTEL added 6 commits June 3, 2026 07:03

Apply suggestion from @bkowalskiINTEL

56176ef

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Apply suggestions from code review

690bb51

Co-authored-by: Bartosz Kowalski <bartosz.kowalski@intel.com> Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Apply suggestion from @bkowalskiINTEL

eb4c879

Fixes missing import Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Remove redundant op support checks in static/dynamic algorithms

c45031b

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Modify docstring for static_quantize

f4119e7

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Comments cleanup

5301022

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

bkowalskiINTEL requested review from anko-intel and Copilot June 3, 2026 14:13

Restore gemma example

c79ef1e

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Copilot AI reviewed Jun 3, 2026

View reviewed changes

bkowalskiINTEL marked this pull request as ready for review June 3, 2026 15:02

anko-intel requested changes Jul 6, 2026

View reviewed changes

bkowalskiINTEL requested review from anko-intel and Copilot July 6, 2026 15:34

Copilot AI reviewed Jul 6, 2026

View reviewed changes

Comment thread neural_compressor/jax/quantization/config.py

Comment thread neural_compressor/jax/quantization/saving.py Outdated

Comment thread test/jax/conftest.py

Comment thread neural_compressor/jax/quantization/saving.py

Copilot started reviewing on behalf of bkowalskiINTEL July 6, 2026 17:09 View session

bkowalskiINTEL and others added 7 commits July 7, 2026 04:21

Fix comments

9d9c3cb

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Fix comments

955f19d

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Apply suggestion from code review

af509c0

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Fix and refactor keras and keras_hub layers registrations

ec72714

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Apply suggestions from code review

19619a4

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Fix config order during model loading

99583fa

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Add unit tests for composable config and layer filtering

925a723

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

bkowalskiINTEL force-pushed the dev/bkowalsk/jax_composable_configs branch from 20ce594 to 96a2397 Compare July 7, 2026 11:26

bkowalskiINTEL added 6 commits July 7, 2026 04:29

Merge branch 'main' into dev/bkowalsk/jax_composable_configs

20e250f

Remove unnecessary comments

bd078eb

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Refactor to use white_list instead of include in per-layer filtering

66801be

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Update tests to use white_list instead of include in per-layer filtering

96a2397

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Unify shared quant config logic + cleanup

350c709

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Align tests tolatest changes

3314eb0

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

bkowalskiINTEL requested a review from Copilot July 8, 2026 10:19

Copilot AI reviewed Jul 8, 2026

View reviewed changes

wpietka reviewed Jul 8, 2026

View reviewed changes

bkowalskiINTEL commented Jul 8, 2026

View reviewed changes

Copilot started reviewing on behalf of bkowalskiINTEL July 8, 2026 13:44 View session

bkowalskiINTEL added 3 commits July 9, 2026 02:28

Add tests for per-class filtering

2230a25

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Apply suggestions from code review + minor config inheritance refactor

767aa25

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Raise error when calib function isn't provided for static quantization

857d4fb

Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>

Uh oh!

Conversation

bkowalskiINTEL commented May 27, 2026

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

anko-intel left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants