Struct QatOptimizer 

pub struct QatOptimizer { /* private fields */ }

QAT optimizer wrapper maintaining full-precision master weights.

The optimizer owns the Adam moment state and a cached fake-quantized view of the master weights, but the FP32 master itself is owned by the caller and passed into Self::step. Each step updates the master in place and re-derives the quantized view from it.

Implementations

impl QatOptimizer

pub fn new( master: &Array1<f64>, config: QatConfig, rng: &mut impl Rng, ) -> Result<Self, GpuOptimError>

Create an optimizer for the given FP32 master weights.

The initial quantized view is derived immediately from master.

Errors

Returns GpuOptimError::InvalidState for empty master weights and propagates any calibration error from the initial quantization.

pub fn quantized_weights(&self) -> &Array1<f64>

The current fake-quantized view of the master weights (what the forward pass uses).

pub fn quant_params(&self) -> Option<&QuantParams>

The integer quantization parameters currently in effect, if the target is integer (always None for fp8 targets).

pub fn step_count(&self) -> u64

Number of Self::step calls performed so far.

pub fn step( &mut self, master: &mut Array1<f64>, grad: &Array1<f64>, rng: &mut impl Rng, ) -> Result<(), GpuOptimError>

Perform one optimizer step: an Adam/AdamW update of the FP32 master followed by re-derivation of the quantized view.

The master weights stay in full precision; only the cached quantized view (see Self::quantized_weights) is rounded onto the grid.

Errors

Returns GpuOptimError::DimensionMismatch if master and grad (or the optimizer’s internal state) disagree on length, and propagates any re-calibration error.

Trait Implementations

impl Clone for QatOptimizer

fn clone(&self) -> QatOptimizer

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for QatOptimizer

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.