burn_mamba::mamba3::moments::recalculated

Trait Mamba3MomentsBackendExt

pub trait Mamba3MomentsBackendExt: Backend {
    // Provided method
    fn mamba3_state_moments_phys(
        xhat_bnlchp: FloatTensor<Self>,
        bhat_bnlchr: FloatTensor<Self>,
        da_bnlh: FloatTensor<Self>,
        rot: FloatTensor<Self>,
        initial_state_bhpr: FloatTensor<Self>,
        valid_len: usize,
        quaternion: bool,
        rope_dim: usize,
        rotate_pairwise: bool,
    ) -> (FloatTensor<Self>, FloatTensor<Self>) { ... }
}

Expand description

Extends the backend with the physical-frame state-moments computation.

The default body runs the serial chunkwise forward on primitive tensors. The Autodiff wrapper overrides it with the memory-efficient recompute backward (in super::backward) that saves only the leaf inputs.

Provided Methods§

Source

fn mamba3_state_moments_phys( xhat_bnlchp: FloatTensor<Self>, bhat_bnlchr: FloatTensor<Self>, da_bnlh: FloatTensor<Self>, rot: FloatTensor<Self>, initial_state_bhpr: FloatTensor<Self>, valid_len: usize, quaternion: bool, rope_dim: usize, rotate_pairwise: bool, ) -> (FloatTensor<Self>, FloatTensor<Self>)

Pooled physical-frame state moments of the combined injections (see mamba3/moments.rs for the math). Returns (m2_bhrr, m1_bhr) raw sums, masked to valid_len.

rot is rank-erased: cumulative angles [b, s, h, a] when quaternion = false (with rope_dim / rotate_pairwise selecting the pairing), cumulative unit quaternions [b, s, h, J, 4] otherwise (rope_dim / rotate_pairwise ignored).