[JAX] Add AbsMaxObserver

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):

test/jax/test_observers.py

@@ -0,0 +1,118 @@
+#!/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 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."""
+    observer = AbsMaxObserver(dtype="float32")
+    inputs = jnp.array([[1.0, -9.0, 2.0]], dtype=jnp.float32)
+    mask = jnp.array([[True, False, True]])
+    observer(inputs, mask=mask)
+    calibrated_range = observer.get_calibrated_range()
+    # The masked -9.0 must be ignored, so max abs is 2.0
+    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)
+
+
+@pytest.mark.parametrize(
+    "activation_dtype,asymmetric,expected",
+    [
+        (jnp.dtype("float8_e4m3fn"), False, AbsMaxObserver),
+        (jnp.dtype("int8"), False, AbsMaxObserver),
+        (jnp.dtype("int8"), True, MinMaxObserver),
+    ],
+)
+def test_static_qdq_layer_uses_expected_observer(activation_dtype, asymmetric, expected):
+    """StaticQDQLayer attaches the observer that matches its quantization scheme."""
+    layer = StaticQDQLayer(
+        name="static_qdq",
+        activation_dtype=activation_dtype,
+        dtype="float32",
+        asymmetric=asymmetric,
+    )
+    layer.add_observers()
+    assert isinstance(layer.input_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")