Enum MaybeQuantized

pub enum MaybeQuantized<M>
where
    M: Quantizable,
{
    Original(M),
    Quantized(M::Quantized),
}

A wrapper for a quantizable module.

Variants

Original(M)

The original module.

Quantized(M::Quantized)

The quantized version of the module.

Implementations

impl<M> MaybeQuantized<M>

where M: Quantizable,

pub fn new(module: M) -> Self

Create a new MaybeQuantized from the original module.

pub fn quantize_with( self, op: impl FnOnce(M) -> Result<M::Quantized, M::QuantizationError>, ) -> Result<Self, M::QuantizationError>

Quantize the module with a custom quantization function.

This is useful if one would like to quantize with a custom group size or bit width.

pub fn is_quantized(&self) -> bool

Check if the module is quantized.

Trait Implementations

impl<M> Clone for MaybeQuantized<M>

where M: Quantizable + Clone, M::Quantized: Clone,

fn clone(&self) -> MaybeQuantized<M>

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<M> Debug for MaybeQuantized<M>

where M: Quantizable + Debug, M::Quantized: Debug,

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

Formats the value using the given formatter.

impl<M, Input> Module<Input> for MaybeQuantized<M>

where M: Quantizable + Module<Input>, M::Quantized: Module<Input, Output = <M as Module<Input>>::Output, Error = <M as Module<Input>>::Error>,

type Output = <M as Module<Input>>::Output

Output type of the module.

type Error = <M as Module<Input>>::Error

Error type for the module.

fn forward(&mut self, x: Input) -> Result<Self::Output, Self::Error>

Forward pass of the module.

fn training_mode(&mut self, mode: bool)

Set whether the module is in training mode.

impl<M> ModuleParameters for MaybeQuantized<M>

where M: Quantizable + ModuleParameters, M::Quantized: ModuleParameters,

fn parameters(&self) -> ModuleParamRef<'_>

Get references to the module parameters.

fn parameters_mut(&mut self) -> ModuleParamMut<'_>

Get mutable references to the module parameters.

fn trainable_parameters(&self) -> ModuleParamRef<'_>

Get references to the trainable parameters. A parameter is trainable if it is NOT frozen.

fn freeze_parameters(&mut self, recursive: bool)

Freeze all parameters in the module.

fn unfreeze_parameters(&mut self, recursive: bool)

Unfreeze all parameters in the module.

fn all_frozen(&self) -> Option<bool>

Check if all parameters in the module are frozen. Returns None if there are no parameters.

fn any_frozen(&self) -> Option<bool>

Check if any parameter in the module is frozen. Returns None if there are no parameters.

fn update(&mut self, parameters: ModuleParam)

Update the module parameters.

fn update_flattened(&mut self, flattened_parameters: FlattenedModuleParam)

Update the module parameters from a flattened representation.

impl<M> Quantizable for MaybeQuantized<M>

where M: Quantizable,

type Quantized = MaybeQuantized<M>

The quantized type.

type QuantizationError = <M as Quantizable>::QuantizationError

The error type for quantization.

fn try_into_quantized( self, group_size: i32, bits: i32, ) -> Result<Self, Self::QuantizationError>

Quantize the module with the specified group size and number of bits.

const DEFAULT_GROUP_SIZE: i32 = 64i32

The default group size for quantization.

const DEFAULT_BITS: i32 = 4i32

The default number of bits for quantization.