burn_mamba/utils/
mod.rs

1//! # Shared utilities
2//!
3//! Building blocks reused across the Mamba families: custom activations and
4//! norms (often fp16-stable variants Burn lacks), loss functions, the
5//! `segsum` / `gqa` tensor helpers, the custom-backward plumbing
6//! (`backend_macros` / `combined_grad` / `primitive`), runtime `sanity` guards,
7//! LR `scheduler`s, and the per-dtype numerical constants below.
8
9use burn::prelude::ToElement;
10use burn::tensor::DType;
11
12/// Macros emitting per-backend `BackendExt` impls + autodiff marker traits.
13#[macro_use]
14pub mod backend_macros;
15/// Learnable `[CLS]`-style class tokens/latents spliced into the sequence.
16pub mod class;
17/// Flatten/unflatten `(y, final_state)` into one tracked tensor for the custom
18/// backward.
19pub mod combined_grad;
20/// Rank-tagged `FloatTensor` primitive wrapper mirroring the `Tensor` method
21/// API, used by the custom-backward gradient math.
22pub(crate) mod fprim;
23/// Virtual-layer → real-weight index scheduling shared by all families.
24pub mod schedule;
25/// Learning-rate schedulers (cosine-annealing + warmup, constant).
26pub mod scheduler;
27/// `max_abs_diff` + gradient-comparison macros used across the test suites.
28#[cfg(test)]
29pub mod test_helpers;
30
31pub use class::{ClassLatent, ClassToken};
32pub use schedule::{BidiSchedule, Schedule};
33pub use scheduler::{ConstantLr, CosineAnnealingLr, Lr};
34
35/// A small `dtype`-specific epsilon for safe division (`x / (y + eps)`),
36/// returned as `f32`.
37///
38/// The value is chosen per float format as the geometric mean (average in
39/// log10 space) of two reference magnitudes: a scaled function of the format's
40/// minimum exponent and the format's machine epsilon.  This places `eps`
41/// comfortably above the denormal/underflow floor while staying negligible
42/// relative to typical activations, for each of f64/f32/f16/bf16.  The
43/// resulting constants are noted inline.  `dtype` is the runtime float dtype of
44/// the tensor being divided (e.g. `x.dtype()`).  Panics on non-float dtypes.
45pub fn div_eps(dtype: DType) -> f32 {
46    match dtype {
47        // 4.0693917e-16
48        DType::F64 => {
49            let raw_exp = -(-f64::MIN_EXP as f32 * 2.3f32).powf(0.35f32);
50            let eps_exp = (f64::EPSILON as f32).log10();
51            let avg = (raw_exp + eps_exp) / 2f32;
52            10f32.powf(avg)
53        }
54        // 8.1584695e-8
55        DType::F32 | DType::Flex32 => {
56            let raw_exp = -(-f32::MIN_EXP as f32 * 2.3f32).powf(0.35f32);
57            let eps_exp = f32::EPSILON.log10();
58            let avg = (raw_exp + eps_exp) / 2f32;
59            10f32.powf(avg)
60        }
61        // 7.1209995e-4
62        DType::F16 => {
63            let raw_exp = -(-burn::tensor::f16::MIN_EXP.to_f32() * 2.3f32).powf(0.35f32);
64            let eps_exp = burn::tensor::f16::EPSILON.to_f32().log10();
65            let avg = (raw_exp + eps_exp) / 2f32;
66            10f32.powf(avg)
67        }
68        // 2.0885676e-5
69        DType::BF16 => {
70            let raw_exp = -(-burn::tensor::bf16::MIN_EXP.to_f32() * 2.3f32).powf(0.35f32);
71            let eps_exp = burn::tensor::bf16::EPSILON.to_f32().log10();
72            let avg = (raw_exp + eps_exp) / 2f32;
73            10f32.powf(avg)
74        }
75        DType::I64
76        | DType::I32
77        | DType::I16
78        | DType::I8
79        | DType::U64
80        | DType::U32
81        | DType::U16
82        | DType::U8
83        | DType::Bool(_) => {
84            unreachable!()
85        }
86        DType::QFloat(_) => {
87            unimplemented!()
88        }
89    }
90}
burn_mamba/utils/mod.rs

burn_mamba/utils/
mod.rs
