Skip to main content

quantize

Function quantize 

Source
pub fn quantize(
    w: impl AsRef<Array>,
    options: QuantizeOptions,
) -> Result<(Array, Array, Option<Array>), QuantizationError>
Expand description

Quantize weights into packed values, scales, and optional biases.

Biases are present only for QuantizationMode::Affine. Inputs must have at least two dimensions and a last dimension divisible by the group size. Invalid mode/group/bit combinations return QuantizationError::InvalidConfiguration. Shapes and dtypes are checked upstream.

use mlx_rs::{Array, ops::{self, QuantizationMode, QuantizeOptions, DequantizeOptions}};
let weights = Array::ones::<f32>(&[32, 64])?;
let mode = QuantizationMode::Mxfp4;
let (packed, scales, biases) = ops::quantize(
    &weights, QuantizeOptions { mode, ..Default::default() },
)?;
assert!(biases.is_none());
let restored = ops::dequantize(
    &packed, &scales, biases.as_ref(),
    DequantizeOptions { mode, ..Default::default() },
)?;
assert_eq!(restored.shape(), weights.shape());