neural_compressor.jax.algorithms.static

Static quantization algorithm entry point for JAX models.

Functions

static_quantize(→ keras.Model)

Quantize model using Static quantization algorithm.

Module Contents

neural_compressor.jax.algorithms.static.static_quantize(model: keras.Model, configs_mapping: OrderedDict[str | str, OrderedDict[str, neural_compressor.jax.quantization.config.JaxBaseConfig]] | None = None, calib_function: Callable = None) → keras.Model[source]

Quantize model using Static 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.

  • calib_function (Optional[Callable]) – Calibration function used to collect activation statistics.

Returns:

The model with quantized layers.

Return type:

keras.Model