ruda-nn 0.21.62

use core::{convert::Infallible,fmt};
use ruda_model::{module::{Module,ModuleDisplay},tensor::{FrozenAwqOps,FrozenNf4Ops,Tensor,backend::Backend}};
use crate::{Linear,LoRALinear,FrozenNf4Linear,Nf4LoRALinear,FrozenAwqLinear,AwqLoRALinear,
    QuantizedLinear,QuantizedLoRALinear};
use super::{AwqTransformerProjection,AwqGroupedQueryAttention,AwqFeedForward,AwqTransformerBlock,
    AwqTransformerStack,AwqTransformerHead,AwqTransformerModel};
use super::AdaptedProjection;

/// Backend-bound identity for the actual stored projection type.
/// The blanket implementation retains the original module value and public field layout.
pub trait TransformerProjectionStorage<B:Backend>:Module<B> {
    /// The original stored module, without a wrapper or extra persistent field.
    type Stored:Module<B>;
}
impl<B:Backend,P:Module<B>> TransformerProjectionStorage<B> for P {type Stored=P;}
/// Actual projection value while retaining the backend type in generic field metadata.
pub type BackendProjection<B,P> = <P as TransformerProjectionStorage<B>>::Stored;

/// Logical projection geometry independent of its execution extension or storage format.
pub trait TransformerProjectionShape<B:Backend>:Module<B>+ModuleDisplay {
    /// Actual logical `[input,output]` widths, without decoding packed parameters.
    fn dimensions(&self) -> [usize;2];
}

/// Actual caller-selected native projection. Implementations retain all leading axes
/// and return the original failure; no dense substitution or quantizer is implied.
pub trait TransformerProjection<B:Backend>:TransformerProjectionShape<B> {
    /// Actual original native projection error.
    type Error:fmt::Debug;
    /// Apply the original selected projection and its original derivative contract.
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error>;
}

impl<B:Backend> TransformerProjectionShape<B> for AwqTransformerProjection<B> {
    fn dimensions(&self) -> [usize;2] {self.dimensions()}
}
impl<B:FrozenAwqOps> TransformerProjection<B> for AwqTransformerProjection<B> {
    type Error=B::AwqError;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {self.forward(input)}
}
impl<B:Backend> TransformerProjectionShape<B> for FrozenAwqLinear<B> {
    fn dimensions(&self) -> [usize;2] {self.dimensions()}
}
impl<B:FrozenAwqOps> TransformerProjection<B> for FrozenAwqLinear<B> {
    type Error=B::AwqError;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {self.forward(input)}
}
impl<B:Backend> TransformerProjectionShape<B> for AwqLoRALinear<B> {
    fn dimensions(&self) -> [usize;2] {self.base.dimensions()}
}
impl<B:FrozenAwqOps> TransformerProjection<B> for AwqLoRALinear<B> {
    type Error=B::AwqError;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {self.forward(input)}
}
impl<B:Backend> TransformerProjectionShape<B> for Linear<B> {
    fn dimensions(&self) -> [usize;2] {self.weight.val().dims()}
}
impl<B:Backend> TransformerProjection<B> for Linear<B> {
    type Error=Infallible;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {Ok(self.forward(input))}
}
impl<B:Backend> TransformerProjectionShape<B> for LoRALinear<B> {
    fn dimensions(&self) -> [usize;2] {self.base.weight.val().dims()}
}
impl<B:Backend> TransformerProjection<B> for LoRALinear<B> {
    type Error=Infallible;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {Ok(self.forward(input))}
}

impl<B:Backend> TransformerProjectionShape<B> for QuantizedLinear<B> {
    fn dimensions(&self) -> [usize;2] {let [output,input]=self.weight.val().dims();[input,output]}
}
impl<B:Backend> TransformerProjection<B> for QuantizedLinear<B> {
    type Error=Infallible;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {Ok(self.forward(input))}
}
impl<B:Backend> TransformerProjectionShape<B> for QuantizedLoRALinear<B> {
    fn dimensions(&self) -> [usize;2] {self.base.dimensions()}
}
impl<B:Backend> TransformerProjection<B> for QuantizedLoRALinear<B> {
    type Error=Infallible;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {Ok(self.forward(input))}
}

/// Per-role native floating, INT2/4/8, FP4 E2M1 or FP8 E4M3/E5M2 projection storage.
/// Plugs directly into the existing projected attention, FFN, head and model graphs.
#[derive(Module,Debug)]
pub enum QuantizedTransformerProjection<B:Backend> {
    /// Original floating projection and its existing training flags.
    Dense(Linear<B>),
    /// Original floating base and actual trainable adapters.
    LoRA(LoRALinear<B>),
    /// Original generic packed weight/scales and optional frozen bias.
    Quantized(QuantizedLinear<B>),
    /// Original generic packed base with independently stored A/B leaves.
    QuantizedLoRA(QuantizedLoRALinear<B>),
}
impl<B:Backend> TransformerProjectionShape<B> for QuantizedTransformerProjection<B> {
    fn dimensions(&self) -> [usize;2] {
        match self {Self::Dense(layer)=>layer.dimensions(),Self::LoRA(layer)=>layer.dimensions(),
            Self::Quantized(layer)=>layer.dimensions(),Self::QuantizedLoRA(layer)=>layer.dimensions()}
    }
}
impl<B:Backend> TransformerProjection<B> for QuantizedTransformerProjection<B> {
    type Error=Infallible;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {
        Ok(match self {Self::Dense(layer)=>layer.forward(input),Self::LoRA(layer)=>layer.forward(input),
            Self::Quantized(layer)=>layer.forward(input),Self::QuantizedLoRA(layer)=>layer.forward(input)})
    }
}
impl<B:Backend> QuantizedTransformerProjection<B> {
    /// Pack a caller-selected actual floating role; preserve existing LoRA IDs, values,
    /// dropout and scaling. Already packed roles are not decoded or silently requantized.
    pub fn pack_float(self, scheme:&ruda_model::tensor::quantization::QuantScheme,
        calibration_dtype:ruda_model::tensor::FloatDType) -> Self {
        match self {
            Self::Dense(layer)=>Self::Quantized(QuantizedLinear::from_linear(layer,scheme,calibration_dtype)),
            Self::LoRA(layer)=>Self::QuantizedLoRA(QuantizedLoRALinear {
                base:QuantizedLinear::from_linear(layer.base,scheme,calibration_dtype),
                adapter_a:layer.adapter_a,adapter_b:layer.adapter_b,dropout:layer.dropout,scale:layer.scale}),
            _=>panic!("pack_float requires an actual floating base, not an already packed checkpoint"),
        }
    }
}

impl<B:Backend> TransformerProjectionShape<B> for AdaptedProjection<B> {
    fn dimensions(&self) -> [usize;2] {
        match self {Self::Dense(layer)=>layer.weight.val().dims(),Self::LoRA(layer)=>layer.base.weight.val().dims()}
    }
}
impl<B:Backend> TransformerProjection<B> for AdaptedProjection<B> {
    type Error=Infallible;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {Ok(self.forward(input))}
}
impl<B:Backend> TransformerProjectionShape<B> for FrozenNf4Linear<B> {
    fn dimensions(&self) -> [usize;2] {[self.input_features,self.output_features]}
}
impl<B:FrozenNf4Ops> TransformerProjection<B> for FrozenNf4Linear<B> {
    type Error=B::Nf4Error;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {self.forward(input)}
}
impl<B:Backend> TransformerProjectionShape<B> for Nf4LoRALinear<B> {
    fn dimensions(&self) -> [usize;2] {[self.base.input_features,self.base.output_features]}
}
impl<B:FrozenNf4Ops> TransformerProjection<B> for Nf4LoRALinear<B> {
    type Error=B::Nf4Error;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {self.forward(input)}
}

/// Explicit dense, dense-LoRA, original NF4 or NF4-LoRA projection.
/// Unlike a mixed AWQ/NF4 graph, this requires only the native NF4 extension.
#[derive(Module,Debug)]
pub enum Nf4TransformerProjection<B:Backend> {
    /// Actual original dense module, including its frozen/trainable flags.
    Dense(Linear<B>),
    /// Actual original floating-base LoRA module.
    LoRA(LoRALinear<B>),
    /// Original immutable high-nibble-first U8 base and FP32 scales/codebook.
    Nf4(FrozenNf4Linear<B>),
    /// Original packed base and independently trainable floating adapters.
    Nf4LoRA(Nf4LoRALinear<B>),
}
impl<B:Backend> TransformerProjectionShape<B> for Nf4TransformerProjection<B> {
    fn dimensions(&self) -> [usize;2] {
        match self {Self::Dense(layer)=>layer.dimensions(),Self::LoRA(layer)=>layer.dimensions(),
            Self::Nf4(layer)=>layer.dimensions(),Self::Nf4LoRA(layer)=>layer.dimensions()}
    }
}
impl<B:FrozenNf4Ops> TransformerProjection<B> for Nf4TransformerProjection<B> {
    type Error=B::Nf4Error;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {
        match self {Self::Dense(layer)=>Ok(layer.forward(input)),Self::LoRA(layer)=>Ok(layer.forward(input)),
            Self::Nf4(layer)=>layer.forward(input),Self::Nf4LoRA(layer)=>layer.forward(input)}
    }
}

/// Original failure retaining which real packed representation was executed.
#[derive(Debug)]
pub enum MixedProjectionError<A:fmt::Debug,N:fmt::Debug> {
    /// Native AWQ execution or derivative failure.
    Awq(A),
    /// Native NF4 execution or derivative failure.
    Nf4(N),
}
impl<A:fmt::Debug,N:fmt::Debug> fmt::Display for MixedProjectionError<A,N> {
    fn fmt(&self,f:&mut fmt::Formatter<'_>) -> fmt::Result {
        match self {Self::Awq(error)=>write!(f,"AWQ projection: {error:?}"),Self::Nf4(error)=>write!(f,"NF4 projection: {error:?}")}
    }
}
impl<A:fmt::Debug,N:fmt::Debug> core::error::Error for MixedProjectionError<A,N> {}

/// Explicit per-role mixed native storage; selecting a role never converts its base.
#[derive(Module,Debug)]
pub enum MixedTransformerProjection<B:Backend> {
    /// Actual original floating values and trainability.
    Dense(Linear<B>),
    /// Actual original floating base and floating adapters.
    LoRA(LoRALinear<B>),
    /// Actual loaded AWQ words, zero points and original scale storage.
    Awq(FrozenAwqLinear<B>),
    /// Actual loaded AWQ base and independent floating adapters.
    AwqLoRA(AwqLoRALinear<B>),
    /// Actual loaded NF4 bytes and original FP32 quantization metadata.
    Nf4(FrozenNf4Linear<B>),
    /// Actual loaded NF4 base and independent floating adapters.
    Nf4LoRA(Nf4LoRALinear<B>),
}
impl<B:Backend> TransformerProjectionShape<B> for MixedTransformerProjection<B> {
    fn dimensions(&self) -> [usize;2] {
        match self {Self::Dense(layer)=>layer.dimensions(),Self::LoRA(layer)=>layer.dimensions(),
            Self::Awq(layer)=>layer.dimensions(),Self::AwqLoRA(layer)=>layer.base.dimensions(),
            Self::Nf4(layer)=>layer.dimensions(),Self::Nf4LoRA(layer)=>layer.dimensions()}
    }
}
impl<B:FrozenAwqOps+FrozenNf4Ops> TransformerProjection<B> for MixedTransformerProjection<B> {
    type Error=MixedProjectionError<B::AwqError,B::Nf4Error>;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {
        match self {Self::Dense(layer)=>Ok(layer.forward(input)),Self::LoRA(layer)=>Ok(layer.forward(input)),
            Self::Awq(layer)=>layer.forward(input).map_err(MixedProjectionError::Awq),
            Self::AwqLoRA(layer)=>layer.forward(input).map_err(MixedProjectionError::Awq),
            Self::Nf4(layer)=>layer.forward(input).map_err(MixedProjectionError::Nf4),
            Self::Nf4LoRA(layer)=>layer.forward(input).map_err(MixedProjectionError::Nf4)}
    }
}

/// One native model graph may mix all existing packed formats without changing any
/// role's original codes, scale storage, block geometry, bias or adapter identities.
#[derive(Module,Debug)]
pub enum UniversalTransformerProjection<B:Backend> {
    /// Existing dense/LoRA/AWQ/NF4 implementations and their original error contracts.
    Mixed(MixedTransformerProjection<B>),
    /// Native generic INT2/4/8 or FP4/FP8 roles, with or without actual LoRA adapters.
    Generic(QuantizedTransformerProjection<B>),
}
impl<B:Backend> From<MixedTransformerProjection<B>> for UniversalTransformerProjection<B> {
    fn from(value:MixedTransformerProjection<B>) -> Self {Self::Mixed(value)}
}
impl<B:Backend> From<QuantizedTransformerProjection<B>> for UniversalTransformerProjection<B> {
    fn from(value:QuantizedTransformerProjection<B>) -> Self {Self::Generic(value)}
}
impl<B:Backend> TransformerProjectionShape<B> for UniversalTransformerProjection<B> {
    fn dimensions(&self) -> [usize;2] {
        match self {Self::Mixed(layer)=>layer.dimensions(),Self::Generic(layer)=>layer.dimensions()}
    }
}
impl<B:FrozenAwqOps+FrozenNf4Ops> TransformerProjection<B> for UniversalTransformerProjection<B> {
    type Error=MixedProjectionError<B::AwqError,B::Nf4Error>;
    fn forward<const D:usize>(&self,input:Tensor<B,D>) -> Result<Tensor<B,D>,Self::Error> {
        match self {Self::Mixed(layer)=>layer.forward(input),
            Self::Generic(layer)=>layer.forward(input).map_err(|never| match never {})}
    }
}

/// Shared native GQA/MQA/self/cross attention graph over explicit projection storage.
pub type ProjectedGroupedQueryAttention<B,P> = AwqGroupedQueryAttention<B,P>;
/// Shared ordinary/gated FFN graph over explicit projection storage.
pub type ProjectedFeedForward<B,P> = AwqFeedForward<B,P>;
/// Shared native dense/packed/cached transformer block graph.
pub type ProjectedTransformerBlock<B,P> = AwqTransformerBlock<B,P>;
/// Shared original ordered block graph with native cache topology.
pub type ProjectedTransformerStack<B,P> = AwqTransformerStack<B,P>;
/// Shared native normalized/dropout output head.
pub type ProjectedTransformerHead<B,P> = AwqTransformerHead<B,P>;
/// Shared complete original embeddings/backbone/final norm/head graph.
pub type ProjectedTransformerModel<B,P> = AwqTransformerModel<B,P>;
/// Native NF4 graph requiring no AWQ backend capability.
pub type Nf4TransformerModel<B> = ProjectedTransformerModel<B,Nf4TransformerProjection<B>>;
/// Native independently selected dense/AWQ/NF4 graph.
pub type MixedTransformerModel<B> = ProjectedTransformerModel<B,MixedTransformerProjection<B>>;