neural_compressor.jax.algorithms.dynamic

Dynamic quantization algorithm entry point for JAX models.

Functions

dynamic_quantize(→ Any)

Quantize model using Dynamic quantization algorithm.

Module Contents

neural_compressor.jax.algorithms.dynamic.dynamic_quantize(model: keras.Model, configs_mapping: OrderedDict[str | str, OrderedDict[str, neural_compressor.common.base_config.BaseConfig]] | None = None, quant_config: neural_compressor.common.base_config.BaseConfig | None = None, *args: Any, **kwargs: Any) → Any[source]

Quantize model using Dynamic quantization algorithm.

Parameters:

  • model (keras.Model) – JAX model to be quantized.

  • configs_mapping (Optional[OrderedDict[Union[str, str], OrderedDict[str, BaseConfig]]]) – Mapping of configurations for the algorithm.

  • quant_config (Optional[BaseConfig]) – Quantization configuration for wrapper selection.

  • *args (Any) – Additional positional arguments (unused).

  • **kwargs (Any) – Additional keyword arguments (unused).

Returns:

The quantized model wrapped for inference.

Return type:

keras.Model