burn_mamba/mamba3/quat_scan/
mod.rs

1//! # Memory-efficient quaternion cumulative-product scan (custom backward)
2//!
3//! The Quaternion4D forward composes its per-step rotations with a cumulative
4//! product over the sequence ([`crate::mamba3::rotation::quat_cumprod`], a
5//! Hillis–Steele parallel scan). Expressed in plain autodiff that scan retains
6//! `O(log seq)` *full-sequence* intermediates for the backward pass — fast, but
7//! memory-hungry.
8//!
9//! This module provides the **recompute** alternative, mirroring the SSD
10//! `SerialRecalculated` design: a [`Mamba3QuatScanBackendExt`] trait whose
11//! `Autodiff` impl ([`backward`]) registers a single custom
12//! [`Backward`](burn::backend::autodiff::ops::Backward) node that saves only the
13//! leaf inputs (the per-step quaternions and the carry) and recomputes the scan
14//! during backprop. The gradient is the exact quaternion VJP of the cumulative
15//! product, evaluated with parallel ops (a prefix product + a reverse-cumsum) —
16//! no token loop, so the saved memory does not buy back a slow backward.
17//!
18//! [`quat_cumprod_recalculated`] is the high-level drop-in for `quat_cumprod`
19//! used by the Quaternion4D `forward` path; the plain-autodiff `quat_cumprod`
20//! stays as the verified reference (the tests assert the two agree on values
21//! **and** gradients).
22
23/// The backend-extension trait, its primitive default body, the per-backend
24/// impls, and the [`quat_cumprod_recalculated`] high-level wrapper.
25pub mod quat_scan;
26
27/// The registered custom `Backward` node (autodiff op) + the recompute gradient
28/// math.
29#[cfg(feature = "autodiff")]
30pub mod backward;
31
32pub use quat_scan::{Mamba3QuatScanBackendExt, quat_cumprod_recalculated};
33
34#[cfg(all(test, feature = "_dev-test"))]
35mod tests;
burn_mamba/mamba3/quat_scan/mod.rs

burn_mamba/mamba3/quat_scan/
mod.rs
