diff --git a/neural_compressor/jax/quantization/layers_static.py b/neural_compressor/jax/quantization/layers_static.py index 5e5b096a3cf..993d93ec131 100644 --- a/neural_compressor/jax/quantization/layers_static.py +++ b/neural_compressor/jax/quantization/layers_static.py @@ -153,6 +153,99 @@ def is_calibrated(self): return not (jnp.isinf(self.min_val.value).any()) +class AbsMaxObserver(keras.layers.Layer): + """Observer that tracks the running maximum absolute value for calibration.""" + + def __init__(self, *args, **kwargs): + """Initialize the absolute-max observer layer. + + Args: + *args: Positional arguments for the base layer. + **kwargs: Keyword arguments for the base layer. + + Returns: + None: Initializes the observer layer. + """ + super().__init__(*args, **kwargs, name="abs_max") + # Track running maximum absolute value as a non-trainable weight + self.max_abs_val = self.add_weight( + shape=(), + initializer=keras.initializers.Constant(-np.inf), + trainable=False, + name="max_abs_val", + dtype=self.compute_dtype, + ) + self.supports_masking = True + + def call(self, inputs, mask=None): + """Update the maximum absolute value statistic during calibration. + + Args: + inputs (jnp.ndarray): Input tensor to observe. + mask (Optional[jnp.ndarray]): Optional mask to ignore padded elements. + + Returns: + jnp.ndarray: The original inputs for passthrough. + """ + if 0 not in inputs.shape: + if mask is not None: + # Expand mask to match input dimensions if needed + if len(mask.shape) < len(inputs.shape): + for _ in range(len(inputs.shape) - len(mask.shape)): + mask = ops.expand_dims(mask, axis=-1) + # Apply mask to exclude masked positions + masked_inputs = ops.where(mask, ops.abs(inputs), jnp.array(float("-inf"), dtype=inputs.dtype)) + batch_max_abs = keras.ops.max(masked_inputs) + else: + batch_max_abs = keras.ops.max(ops.abs(inputs)) + + self.max_abs_val.assign(keras.ops.maximum(self.max_abs_val, batch_max_abs)) + return inputs + + def build(self, input_shape): + """Override build with no additional variables. + + Args: + input_shape (Tuple[int, ...]): Input shape for the layer. + + Returns: + None: No additional variables are created. + """ + pass + + def get_calibrated_range(self): + """Return the calibrated maximum absolute value. + + Returns: + jnp.ndarray: Tensor containing the maximum absolute value. + """ + return ops.array((self.max_abs_val,)) + + def is_calibrated(self): + """Check if the observer has valid calibration data. + + Returns: + bool: True if calibrated, False if the max abs value is still at its initial value. + """ + return not (jnp.isinf(self.max_abs_val.value).any()) + + +def get_activation_observer(activation_dtype, asymmetric, dtype_policy): + """Select the appropriate activation observer for a quantization scheme. + + Args: + activation_dtype (jnp.dtype): Activation dtype used for quantization. + asymmetric (bool): Whether asymmetric quantization is used. + dtype_policy (keras.DTypePolicy): dtype policy for the observer layer. + + Returns: + keras.layers.Layer: An instance of MinMaxObserver or AbsMaxObserver. + """ + if asymmetric and jnp.issubdtype(activation_dtype, jnp.integer): + return MinMaxObserver(dtype=dtype_policy) + return AbsMaxObserver(dtype=dtype_policy) + + class StaticQDQLayer(SaveableLayerMixin, keras.layers.Layer): """Layer that applies static quantize-dequantize to activations.""" @@ -208,7 +301,7 @@ def add_observers(self): if self.fixed_range is not None: return self._tracker.unlock() - self.input_observer = MinMaxObserver(dtype=self.dtype_policy) + self.input_observer = get_activation_observer(self.activation_dtype, self._is_asymmetric, self.dtype_policy) self._tracker.lock() def add_variables(self): @@ -487,7 +580,7 @@ def add_observers(self): None: Adds observer layers. """ self._tracker.unlock() - self.input_observer = MinMaxObserver(dtype=self.dtype_policy) + self.input_observer = get_activation_observer(self.activation_dtype, self._is_int8, self.dtype_policy) self._tracker.lock() def add_variables(self): diff --git a/test/jax/test_observers.py b/test/jax/test_observers.py new file mode 100644 index 00000000000..6dd8fb7a691 --- /dev/null +++ b/test/jax/test_observers.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# Copyright (c) 2026 Intel Corporation +# +# 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 activation observers used during static quantization calibration. + +Verifies that AbsMaxObserver tracks the maximum absolute value, and that eligible +layers select the correct observer: MinMaxObserver only for asymmetric integer +quantization, and AbsMaxObserver otherwise (fp8 and symmetric int8). +""" + +import pytest +from jax import numpy as jnp +from keras.src import backend + +from neural_compressor.jax.quantization.layers_static import ( + AbsMaxObserver, + MinMaxObserver, + StaticQDQLayer, + get_activation_observer, +) + +# Mark all tests in this file as smoke tests +pytestmark = pytest.mark.smoke_test + + +def test_abs_max_observer_tracks_max_abs(): + """AbsMaxObserver records the running maximum absolute value across calls.""" + observer = AbsMaxObserver(dtype="float32") + + assert not observer.is_calibrated() + + observer(jnp.array([1.0, -2.0, 0.5], dtype=jnp.float32)) + observer(jnp.array([-5.0, 3.0], dtype=jnp.float32)) + observer(jnp.array([4.0, -1.0], dtype=jnp.float32)) + + assert observer.is_calibrated() + + calibrated_range = observer.get_calibrated_range() + # Only the maximum absolute value is returned + assert float(calibrated_range[0]) == pytest.approx(5.0) + + +def test_abs_max_observer_passthrough(): + """AbsMaxObserver returns its inputs unchanged.""" + observer = AbsMaxObserver(dtype="float32") + inputs = jnp.array([1.0, -2.0, 3.0], dtype=jnp.float32) + outputs = observer(inputs) + assert jnp.array_equal(inputs, outputs) + + +def test_abs_max_observer_respects_mask(): + """AbsMaxObserver ignores masked-out positions.""" + + def _get_fields(): + observer = AbsMaxObserver(dtype="float32") + inputs = jnp.array([[1.0, -9.0, 2.0]], dtype=jnp.float32) + mask = jnp.array([[True, False, True]]) + return observer, inputs, mask + + observer, inputs, mask = _get_fields() + observer(inputs, mask=mask) + + calibrated_range = observer.get_calibrated_range() + assert float(calibrated_range[0]) == pytest.approx(2.0) + + observer, inputs, mask = _get_fields() + backend.set_keras_mask(inputs, mask) + observer(inputs) + + calibrated_range = observer.get_calibrated_range() + assert float(calibrated_range[0]) == pytest.approx(2.0) + + +@pytest.mark.parametrize( + "activation_dtype,asymmetric,expected", + [ + (jnp.dtype("float8_e4m3fn"), False, AbsMaxObserver), + (jnp.dtype("float8_e5m2"), False, AbsMaxObserver), + (jnp.dtype("float8_e4m3fn"), True, AbsMaxObserver), # fp8 is always symmetric + (jnp.dtype("int8"), False, AbsMaxObserver), # symmetric int8 + (jnp.dtype("int8"), True, MinMaxObserver), # asymmetric int8 + ], +) +def test_get_activation_observer_selection(activation_dtype, asymmetric, expected): + """The helper selects MinMaxObserver only for asymmetric integer quantization.""" + observer = get_activation_observer(activation_dtype, asymmetric, dtype_policy="float32") + assert isinstance(observer, expected) + + +def test_static_qdq_layer_no_observer_with_fixed_range(): + """No observer is attached when a fixed range is provided.""" + layer = StaticQDQLayer( + name="static_qdq_fixed", + activation_dtype=jnp.dtype("float8_e4m3fn"), + dtype="float32", + fixed_range=(-3.0, 3.0), + ) + layer.add_observers() + assert not hasattr(layer, "input_observer")