neural_compressor.jax.quantization.layers_static

Static quantized layer implementations for JAX-backed Keras models.

Classes

MinMaxObserver	Observer that tracks running min/max values for calibration.
AbsMaxObserver	Observer that tracks the running maximum absolute value for calibration.
StaticQDQLayer	Layer that applies static quantize-dequantize to activations.
QStaticDenseMixin	Mixin that adds static quantization to dense-like layers.
QStaticDense	Statically quantized Dense layer.
QStaticEinsumDense	Statically quantized EinsumDense layer.
QStaticConv2DMixin	Mixin that adds static quantization to Conv2D layers.
QStaticConv2d	Statically quantized Conv2D layer.
QStaticMultiHeadAttention	Statically quantized MultiHeadAttention layer.
QStaticCachedGemma3Attention	Statically quantized CachedGemma3Attention layer.
QStaticGemma3VisionAttention	Statically quantized Gemma3VisionAttention layer.
QStaticReversibleEmbedding	Statically quantized ReversibleEmbedding layer.

Functions

register_static_quantized_layer(clso)	Register quantized layer class for an original layer class.
get_activation_observer(activation_dtype, asymmetric, ...)	Select the appropriate activation observer for a quantization scheme.

Module Contents

neural_compressor.jax.quantization.layers_static.register_static_quantized_layer(clso)

Register quantized layer class for an original layer class.

Parameters:

clso (type) – Original layer class to map to a quantized implementation.

Returns:

Decorator that registers the quantized class.

Return type:

Callable

class neural_compressor.jax.quantization.layers_static.MinMaxObserver(*args, **kwargs)

Observer that tracks running min/max values for calibration.

class neural_compressor.jax.quantization.layers_static.AbsMaxObserver(*args, **kwargs)

Observer that tracks the running maximum absolute value for calibration.

neural_compressor.jax.quantization.layers_static.get_activation_observer(activation_dtype, asymmetric, dtype_policy)

Select the appropriate activation observer for a quantization scheme.

Parameters:

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

An instance of MinMaxObserver or AbsMaxObserver.

Return type:

keras.layers.Layer

class neural_compressor.jax.quantization.layers_static.StaticQDQLayer(name, activation_dtype, dtype='float32', asymmetric=False, const_scale=False, fixed_range=None)

Layer that applies static quantize-dequantize to activations.

class neural_compressor.jax.quantization.layers_static.QStaticDenseMixin

Mixin that adds static quantization to dense-like layers.

class neural_compressor.jax.quantization.layers_static.QStaticDense

Statically quantized Dense layer.

class neural_compressor.jax.quantization.layers_static.QStaticEinsumDense

Statically quantized EinsumDense layer.

class neural_compressor.jax.quantization.layers_static.QStaticConv2DMixin

Mixin that adds static quantization to Conv2D layers.

class neural_compressor.jax.quantization.layers_static.QStaticConv2d

Statically quantized Conv2D layer.

class neural_compressor.jax.quantization.layers_static.QStaticMultiHeadAttention

Statically quantized MultiHeadAttention layer.

class neural_compressor.jax.quantization.layers_static.QStaticCachedGemma3Attention

Statically quantized CachedGemma3Attention layer.

class neural_compressor.jax.quantization.layers_static.QStaticGemma3VisionAttention

Statically quantized Gemma3VisionAttention layer.

class neural_compressor.jax.quantization.layers_static.QStaticReversibleEmbedding

Statically quantized ReversibleEmbedding layer.