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());