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.
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.
QStaticMultiHeadAttention	Statically quantized MultiHeadAttention layer.
QStaticCachedGemma3Attention	Statically quantized CachedGemma3Attention layer.
QStaticGemma3VisionAttention	Statically quantized Gemma3VisionAttention layer.
QStaticRotaryEmbedding	Statically quantized RotaryEmbedding layer.
QStaticReversibleEmbedding	Statically quantized ReversibleEmbedding layer.

Functions

register_static_quantized_layer(clso)

Register quantized layer class for an original layer class.

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.StaticQDQLayer(name, activation_dtype, asymmetric=False)

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.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.QStaticRotaryEmbedding

Statically quantized RotaryEmbedding layer.

class neural_compressor.jax.quantization.layers_static.QStaticReversibleEmbedding

Statically quantized ReversibleEmbedding layer.