neural_compressor.jax.quantization.layers_dynamic

Dynamic quantized layer implementations for JAX-backed Keras models.

Classes

DynamicQDQLayer	Layer that applies dynamic quantize-dequantize to activations.
QDynamicDenseMixin	Mixin that adds dynamic quantization to dense-like layers.
QDynamicDense	Dynamically quantized Dense layer.
QDynamicEinsumDense	Dynamically quantized EinsumDense layer.
QDynamicMultiHeadAttention	Dynamically quantized MultiHeadAttention layer.
QDynamicCachedGemma3Attention	Dynamically quantized CachedGemma3Attention layer.
QDynamicGemma3VisionAttention	Dynamically quantized Gemma3VisionAttention layer.
QDynamicReversibleEmbedding	Dynamically quantized ReversibleEmbedding layer.

Functions

register_dynamic_quantized_layer(clso)

Register quantized layer class for an original layer class.

Module Contents

neural_compressor.jax.quantization.layers_dynamic.register_dynamic_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_dynamic.DynamicQDQLayer(name, activation_dtype, asymmetric=False)

Layer that applies dynamic quantize-dequantize to activations.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicDenseMixin

Mixin that adds dynamic quantization to dense-like layers.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicDense

Dynamically quantized Dense layer.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicEinsumDense

Dynamically quantized EinsumDense layer.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicMultiHeadAttention

Dynamically quantized MultiHeadAttention layer.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicCachedGemma3Attention

Dynamically quantized CachedGemma3Attention layer.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicGemma3VisionAttention

Dynamically quantized Gemma3VisionAttention layer.

class neural_compressor.jax.quantization.layers_dynamic.QDynamicReversibleEmbedding

Dynamically quantized ReversibleEmbedding layer.