Module backward 

The registered custom Backward node (autodiff op) + the recompute gradient math.

Custom autodiff node for the quaternion cumulative-product scan

Implements Mamba3QuatScanBackendExt for Autodiff<B> via a single Burn [Backward] node. The forward keeps only its two leaf inputs (the per-step quaternions q and the carry init); backprop recomputes the prefix product and evaluates the exact quaternion VJP of the cumulative product, so the O(log seq) scan intermediates are never retained.

Gradient math

The scan is cum[t] = qₜ ⊗ cum[t-1], cum[-1] = init (so cum[t] = Pₜ ⊗ init with the prefix product Pₜ = qₜ ⊗ ⋯ ⊗ q₀). Quaternion multiplication is bilinear, and for the real inner product the per-factor VJPs are exact for all quaternions:

  o = a ⊗ b ,  cotangent ḡ   ⟹   grad_a = ḡ ⊗ conj(b) ,  grad_b = conj(a) ⊗ ḡ .

Let G[t] be the total cotangent reaching cum[t] (the upstream d_cum[t] plus what flows back from cum[t+1]). The reverse recurrence G[t] = d_cum[t] + conj(qₜ₊₁) ⊗ G[t+1] telescopes, and for unit quaternions (which the rotation always produces) collapses to a parallel-friendly closed form:

  S[t]   = Σ_{s≥t} conj(Pₛ) ⊗ d_cum[s]        (suffix-sum = reverse cumsum)
  G[t]   = Pₜ ⊗ S[t]
  d_q[t] = G[t] ⊗ conj(cum[t-1])              (cum[-1] = init)
  d_init = S[0]

final_carry = cum[:, −1] is a high-level autodiff slice (in quat_cumprod_recalculated), so its gradient is already folded into d_cum at the last position before this node runs — the node has a single output.

Functions

quat_scan_backward 🔒

Recompute-based gradient of the quaternion cumulative product.