Skip to content

[JAX] Add AbsMaxObserver - #2585

Open
wpietka wants to merge 5 commits into
mainfrom
dev/wpietkax/absmax-observer
Open

wpietka wants to merge 5 commits into
mainfrom
dev/wpietkax/absmax-observer

Conversation

@wpietka

@wpietka wpietka commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

Type of Change

feature

Description

Implements new AbsMaxObserver for the JAX static quantization calibration path and wires it into eligible layers.

Expected Behavior & Potential Risk

  • MinMaxObserver is used for asymmetric int8 quantization
  • AbsMaxObserver is used for every other case

How has this PR been tested?

  • New unit tests: test/jax/test_observers.py

Dependency Change?

None

Signed-off-by: Wojciech Piętka <wojciechx.pietka@intel.com>
Signed-off-by: Wojciech Piętka <wojciechx.pietka@intel.com>

@anko-intel anko-intel left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see also:
max_val = jnp.max(jnp.abs(orig_weight), axis=axis, keepdims=True)
in neural_compressor/jax/utils/utility.py
Maybe we can remove max(abs( there?

Comment thread neural_compressor/jax/quantization/layers_static.py Outdated
Signed-off-by: Wojciech Piętka <wojciechx.pietka@intel.com>
@wpietka

wpietka commented Sep 16, 2026

Copy link
Copy Markdown
Contributor Author

see also: max_val = jnp.max(jnp.abs(orig_weight), axis=axis, keepdims=True) in neural_compressor/jax/utils/utility.py Maybe we can remove max(abs( there?

get_scale() is used both for activations and kernel. While for activations it's redundant due to AbsMaxObserver, it is still necessary for kernel scale

@anko-intel anko-intel left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Signed-off-by: Wojciech Piętka <wojciechx.pietka@intel.com>
Signed-off-by: Wojciech Piętka <wojciechx.pietka@intel.com>

@anko-intel anko-intel left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants