Skip to main content

entrenar/quant/fake_quantize/
ops.rs

1//! Convenience functions for fake quantization operations.
2
3use crate::Tensor;
4
5use super::config::FakeQuantConfig;
6use super::quantize::FakeQuantize;
7
8/// Convenience function for fake quantization forward pass
9pub fn fake_quantize(input: &Tensor, bits: usize, symmetric: bool) -> Tensor {
10    let config = if symmetric {
11        FakeQuantConfig::symmetric(bits)
12    } else {
13        FakeQuantConfig::asymmetric(bits)
14    };
15    let mut fq = FakeQuantize::new(config);
16    fq.forward_with_calibration(input)
17}
18
19/// Convenience function for STE backward pass
20pub fn ste_backward(grad_output: &Tensor) -> Tensor {
21    // STE: gradient passes through unchanged
22    grad_output.clone()
23}