neural_compressor.jax.quantization.saving

Serialization helpers for JAX quantized Keras models.

Classes

VersionManager	Handle version metadata for serialized quantized models.
SaveableLayerMixin	Mixin for saving and loading quantized layer variables.
KerasQuantizedModelBackboneWrapper	Wrapper that preserves quantization config when saving Keras backbones.
KerasQuantizedModelWrapperMixin	Wrapper that preserves quantization config for Keras tasks.
KerasQuantizedModelWrapper	Generic quantized model wrapper for Keras models without specific backbone or task structure.
KerasQuantizedGemmaWrapper	Quantized wrapper for Gemma3CausalLM models.
KerasQuantizedViTWrapper	Quantized wrapper for ViTImageClassifier models.
KerasQuantizedTokenizerWrapper	Quantized wrapper for Gemma3Tokenizer models.

Functions

prepare_deserialized_quantized_model(...)

Transform a loaded quantized model.

Module Contents

class neural_compressor.jax.quantization.saving.VersionManager

Handle version metadata for serialized quantized models.

class neural_compressor.jax.quantization.saving.SaveableLayerMixin

Mixin for saving and loading quantized layer variables.

class neural_compressor.jax.quantization.saving.KerasQuantizedModelBackboneWrapper(model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig | None = None)

Wrapper that preserves quantization config when saving Keras backbones.

class neural_compressor.jax.quantization.saving.KerasQuantizedModelWrapperMixin(model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig | None = None)

Wrapper that preserves quantization config for Keras tasks.

class neural_compressor.jax.quantization.saving.KerasQuantizedModelWrapper(model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig | None = None)

Generic quantized model wrapper for Keras models without specific backbone or task structure.

class neural_compressor.jax.quantization.saving.KerasQuantizedGemmaWrapper(model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig | None = None)

Quantized wrapper for Gemma3CausalLM models.

class neural_compressor.jax.quantization.saving.KerasQuantizedViTWrapper(model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig | None = None)

Quantized wrapper for ViTImageClassifier models.

class neural_compressor.jax.quantization.saving.KerasQuantizedTokenizerWrapper(model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig | None = None)

Quantized wrapper for Gemma3Tokenizer models.

neural_compressor.jax.quantization.saving.prepare_deserialized_quantized_model(model: keras.Model, quant_config: neural_compressor.jax.quantization.config.JaxBaseConfig) → KerasQuantizedModelWrapperMixin | KerasQuantizedModelBackboneWrapper

Transform a loaded quantized model.

It prepares the model for inference by preparing the quantized layers. :param model: Loaded base keras model. :type model: keras.Model :param quant_config: Quantization configuration. :type quant_config: JaxBaseConfig

Returns:

The transformed quantized model/backbone wrapper.

Return type:

Union[KerasQuantizedModelWrapperMixin, KerasQuantizedModelBackboneWrapper]