Function fp8_linear

cake-core/src/utils/fp8.rs:247–263

Create an Fp8Linear from a VarBuilder (no bias).

(
    in_features: usize,
    out_features: usize,
    vb: candle_nn::VarBuilder,
)

/// Create an Fp8Linear from a VarBuilder (no bias).
pub fn fp8_linear(
    in_features: usize,
    out_features: usize,
    vb: candle_nn::VarBuilder,
) -> candle_core::Result<Fp8Linear> {
    let weight = vb.get_unchecked_dtype("weight", candle_core::DType::F32)?;
    if weight.dims().len() == 2 && weight.dim(0)? == out_features && weight.dim(1)? == in_features {
        Ok(Fp8Linear::new(weight, None))
    } else if weight.dims().len() == 2
        && weight.dim(0)? == in_features
        && weight.dim(1)? == out_features
    {
        Ok(Fp8Linear::new(weight.t()?, None))
    } else {
        Ok(Fp8Linear::new(weight, None))
    }
}

/// Create an Fp8Linear from a VarBuilder (with bias).
pub fn fp8_linear_b(