rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)

Browse Source 
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU
backward work. Verified formatting-only via diff sampling; no logic
changed.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
This commit is contained in:
osobh
2026-08-10 07:09:36 -07:00
co-authored by Claude Sonnet 5
parent ad6405663f
commit 4aaa36a57a
305 changed files with 25537 additions and 18337 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/ssl/vjepa.rs
+173 -49
View File
@@ -6,7 +6,10 @@
//! - Temporal context masking strategy
//! - Neuro-JEPA extension: apply same paradigm to EEG/MEG signal segments
use super::jepa::{BlockMaskConfig, JepaPredictor, EmaTargetEncoder, JepaLossResult, jepa_loss, FeatureBank, JepaEvaluator};
use super::jepa::{
BlockMaskConfig, EmaTargetEncoder, FeatureBank, JepaEvaluator, JepaLossResult, JepaPredictor,
jepa_loss,
};
// ============================================================================
// Video Patch Configuration
@@ -102,7 +105,9 @@ impl PatchEmbed3D {
let mut lcg = s.wrapping_mul(2654435761).wrapping_add(1);
let scale = (2.0f32 / size as f32).sqrt();
for x in v.iter_mut() {
lcg = lcg.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
lcg = lcg
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
let u = ((lcg >> 11) as f32) / (1u64 << 53) as f32;
*x = (u * 2.0 - 1.0) * scale;
}
@@ -148,7 +153,10 @@ impl PatchEmbed3D {
for pw in 0..sp {
let col = wi * sp + pw;
for c in 0..channels {
let flat_idx = frame * h * w * channels + row * w * channels + col * channels + c;
let flat_idx = frame * h * w * channels
+ row * w * channels
+ col * channels
+ c;
if flat_idx < video.len() {
patch[pi] = video[flat_idx];
}
@@ -217,9 +225,15 @@ pub struct TubeMaskResult {
}
impl TubeMaskResult {
pub fn num_context(&self) -> usize { self.context_indices.len() }
pub fn num_target_patches(&self) -> usize { self.all_target_indices.len() }
pub fn num_target_tubes(&self) -> usize { self.target_tubes.len() }
pub fn num_context(&self) -> usize {
self.context_indices.len()
}
pub fn num_target_patches(&self) -> usize {
self.all_target_indices.len()
}
pub fn num_target_tubes(&self) -> usize {
self.target_tubes.len()
}
}
/// Video tube masking strategy
@@ -231,16 +245,28 @@ pub struct TubeMaskStrategy {
impl TubeMaskStrategy {
pub fn new(config: TubeMaskConfig) -> Self {
Self { config, lcg: 9876543210987654321 }
Self {
config,
lcg: 9876543210987654321,
}
}
fn lcg_next(&mut self) -> u64 {
self.lcg = self.lcg.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
self.lcg = self
.lcg
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
self.lcg
}
fn rand_f64(&mut self) -> f64 { (self.lcg_next() >> 11) as f64 / (1u64 << 53) as f64 }
fn rand_range(&mut self, lo: f64, hi: f64) -> f64 { lo + self.rand_f64() * (hi - lo) }
fn rand_usize(&mut self, n: usize) -> usize { (self.lcg_next() % n as u64) as usize }
fn rand_f64(&mut self) -> f64 {
(self.lcg_next() >> 11) as f64 / (1u64 << 53) as f64
}
fn rand_range(&mut self, lo: f64, hi: f64) -> f64 {
lo + self.rand_f64() * (hi - lo)
}
fn rand_usize(&mut self, n: usize) -> usize {
(self.lcg_next() % n as u64) as usize
}
/// Sample a spatial block (set of spatial positions) and extend across all T
fn sample_tube(&mut self, video_cfg: &VideoPatchConfig) -> Vec<usize> {
@@ -258,7 +284,7 @@ impl TubeMaskStrategy {
let block_w = (area as f64 / block_h as f64).ceil() as usize;
let block_w = block_w.clamp(1, sg);
let top = self.rand_usize((sg - block_h + 1).max(1));
let top = self.rand_usize((sg - block_h + 1).max(1));
let left = self.rand_usize((sg - block_w + 1).max(1));
let mut indices = Vec::with_capacity(ts * block_h * block_w);
@@ -285,7 +311,9 @@ impl TubeMaskStrategy {
for _ in 0..cfg.num_target_tubes {
let tube = self.sample_tube(video_cfg);
for &idx in &tube { all_target_set.insert(idx); }
for &idx in &tube {
all_target_set.insert(idx);
}
target_tubes.push(tube);
}
let all_target_indices: Vec<usize> = all_target_set.into_iter().collect();
@@ -312,8 +340,14 @@ impl TubeMaskStrategy {
}
/// Batch generation
pub fn generate_batch(&mut self, batch_size: usize, video_cfg: &VideoPatchConfig) -> Vec<TubeMaskResult> {
(0..batch_size).map(|i| self.generate(video_cfg, i as u64)).collect()
pub fn generate_batch(
&mut self,
batch_size: usize,
video_cfg: &VideoPatchConfig,
) -> Vec<TubeMaskResult> {
(0..batch_size)
.map(|i| self.generate(video_cfg, i as u64))
.collect()
}
}
@@ -383,11 +417,17 @@ impl MockVideoEncoder {
let mut lcg = 54321u64;
let weights: Vec<f32> = (0..total_patches * embed_dim)
.map(|_| {
lcg = lcg.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
lcg = lcg
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((lcg >> 11) as f32 / (1u64 << 53) as f32) * 2.0 - 1.0
})
.collect();
Self { embed_dim, total_patches, weights }
Self {
embed_dim,
total_patches,
weights,
}
}
pub fn encode(&self, patch_indices: &[usize]) -> Vec<f32> {
@@ -396,15 +436,23 @@ impl MockVideoEncoder {
let mut out = vec![0.0f32; n * d];
for (i, &pi) in patch_indices.iter().enumerate() {
let pi = pi.min(self.total_patches - 1);
for dd in 0..d { out[i * d + dd] = self.weights[pi * d + dd]; }
for dd in 0..d {
out[i * d + dd] = self.weights[pi * d + dd];
}
}
out
}
pub fn l2_normalize(reps: &mut Vec<f32>, n: usize, d: usize) {
for i in 0..n {
let norm = (0..d).map(|dd| reps[i * d + dd].powi(2)).sum::<f32>().sqrt().max(1e-8);
for dd in 0..d { reps[i * d + dd] /= norm; }
let norm = (0..d)
.map(|dd| reps[i * d + dd].powi(2))
.sum::<f32>()
.sqrt()
.max(1e-8);
for dd in 0..d {
reps[i * d + dd] /= norm;
}
}
}
}
@@ -425,7 +473,7 @@ pub struct VJepaStepMetrics {
pub struct VJepaTrainer {
pub config: VJepaConfig,
pub context_encoder: MockVideoEncoder,
pub target_ema: Vec<f32>, // shadow weights of context_encoder
pub target_ema: Vec<f32>, // shadow weights of context_encoder
pub predictor: JepaPredictor,
pub mask_strategy: TubeMaskStrategy,
pub step: usize,
@@ -441,10 +489,24 @@ impl VJepaTrainer {
let context_encoder = MockVideoEncoder::new(embed_dim, total_patches);
let target_ema = context_encoder.weights.clone();
let predictor = JepaPredictor::new(embed_dim, pred_dim, config.predictor_depth, num_heads, total_patches);
let predictor = JepaPredictor::new(
embed_dim,
pred_dim,
config.predictor_depth,
num_heads,
total_patches,
);
let mask_strategy = TubeMaskStrategy::new(config.mask.clone());
Self { config, context_encoder, target_ema, predictor, mask_strategy, step: 0, total_steps }
Self {
config,
context_encoder,
target_ema,
predictor,
mask_strategy,
step: 0,
total_steps,
}
}
fn tau(&self) -> f64 {
@@ -459,13 +521,19 @@ impl VJepaTrainer {
let mut out = vec![0.0f32; n * d];
for (i, &pi) in patch_indices.iter().enumerate() {
let pi = pi.min(np - 1);
for dd in 0..d { out[i * d + dd] = self.target_ema[pi * d + dd]; }
for dd in 0..d {
out[i * d + dd] = self.target_ema[pi * d + dd];
}
}
out
}
fn ema_update(&mut self, tau: f64) {
for (sw, &ow) in self.target_ema.iter_mut().zip(self.context_encoder.weights.iter()) {
for (sw, &ow) in self
.target_ema
.iter_mut()
.zip(self.context_encoder.weights.iter())
{
*sw = (tau as f32) * *sw + (1.0 - tau as f32) * ow;
}
}
@@ -487,7 +555,9 @@ impl VJepaTrainer {
let _ = bi;
let ctx_reps = self.context_encoder.encode(&mask.context_indices);
let all_target = &mask.all_target_indices;
let predicted = self.predictor.forward(&ctx_reps, &mask.context_indices, all_target);
let predicted = self
.predictor
.forward(&ctx_reps, &mask.context_indices, all_target);
let mut target_reps = self.target_encode(all_target);
let n_tgt = all_target.len();
MockVideoEncoder::l2_normalize(&mut target_reps, n_tgt, embed_dim);
@@ -501,10 +571,18 @@ impl VJepaTrainer {
cursor += tube.len();
}
let loss_result = jepa_loss(&pred_n, &target_reps, embed_dim, &mask.target_tubes, &offsets);
let loss_result = jepa_loss(
&pred_n,
&target_reps,
embed_dim,
&mask.target_tubes,
&offsets,
);
total_loss += loss_result.loss;
for (k, &bl) in loss_result.block_losses.iter().enumerate() {
if k < tube_loss_acc.len() { tube_loss_acc[k] += bl; }
if k < tube_loss_acc.len() {
tube_loss_acc[k] += bl;
}
}
total_ctx += mask.num_context();
total_tgt += mask.num_target_patches();
@@ -551,9 +629,9 @@ pub struct NeuroJepaConfig {
impl Default for NeuroJepaConfig {
fn default() -> Self {
Self {
num_channels: 64, // typical EEG montage
segment_len: 64, // ~250ms at 256Hz
total_samples: 1024, // ~4s window at 256Hz
num_channels: 64, // typical EEG montage
segment_len: 64, // ~250ms at 256Hz
total_samples: 1024, // ~4s window at 256Hz
embed_dim: 256,
mask_ratio: 0.75,
ema_tau: 0.996,
@@ -590,13 +668,20 @@ pub struct NeuroMaskStrategy {
}
impl NeuroMaskStrategy {
pub fn new() -> Self { Self { lcg: 1111111111 } }
pub fn new() -> Self {
Self { lcg: 1111111111 }
}
fn lcg_next(&mut self) -> u64 {
self.lcg = self.lcg.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
self.lcg = self
.lcg
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
self.lcg
}
fn rand_usize(&mut self, n: usize) -> usize { (self.lcg_next() % n as u64) as usize }
fn rand_usize(&mut self, n: usize) -> usize {
(self.lcg_next() % n as u64) as usize
}
/// Generate channel-tube masks: select `mask_fraction * num_channels` channels,
/// mask ALL time segments for those channels (tube masking in channel dimension).
@@ -630,11 +715,18 @@ impl NeuroMaskStrategy {
}
}
}
NeuroMaskResult { context_indices, target_indices }
NeuroMaskResult {
context_indices,
target_indices,
}
}
}
impl Default for NeuroMaskStrategy { fn default() -> Self { Self::new() } }
impl Default for NeuroMaskStrategy {
fn default() -> Self {
Self::new()
}
}
// ============================================================================
// Tests — Batch 23 (V-JEPA + Neuro-JEPA)
@@ -697,7 +789,10 @@ mod tests {
let video: Vec<f32> = (0..2 * 8 * 8).map(|i| i as f32 * 0.01).collect();
let out = embed.embed(&video, 1);
let any_nonzero = out.iter().any(|&v| v.abs() > 1e-8);
assert!(any_nonzero, "non-zero input should produce non-zero embeddings");
assert!(
any_nonzero,
"non-zero input should produce non-zero embeddings"
);
}
// ── Tube masking ──────────────────────────────────────────────────────────
@@ -705,7 +800,11 @@ mod tests {
#[test]
fn test_tube_mask_basic() {
let video_cfg = VJepaConfig::small_test().video;
let mask_cfg = TubeMaskConfig { mask_ratio: 0.75, num_target_tubes: 2, ..Default::default() };
let mask_cfg = TubeMaskConfig {
mask_ratio: 0.75,
num_target_tubes: 2,
..Default::default()
};
let mut strategy = TubeMaskStrategy::new(mask_cfg);
let result = strategy.generate(&video_cfg, 0);
assert!(!result.context_indices.is_empty());
@@ -732,8 +831,12 @@ mod tests {
let total = video_cfg.total_patches();
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
let result = strategy.generate(&video_cfg, 0);
for &ci in &result.context_indices { assert!(ci < total); }
for &ti in &result.all_target_indices { assert!(ti < total); }
for &ci in &result.context_indices {
assert!(ci < total);
}
for &ti in &result.all_target_indices {
assert!(ti < total);
}
}
#[test]
@@ -792,8 +895,13 @@ mod tests {
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
let batch = strategy.generate_batch(8, &video_cfg);
// At least some pairs in the batch should differ
let any_differ = batch.windows(2).any(|w| w[0].all_target_indices != w[1].all_target_indices);
assert!(any_differ, "different seeds should produce different masks in at least one pair");
let any_differ = batch
.windows(2)
.any(|w| w[0].all_target_indices != w[1].all_target_indices);
assert!(
any_differ,
"different seeds should produce different masks in at least one pair"
);
}
// ── V-JEPA trainer ────────────────────────────────────────────────────────
@@ -832,7 +940,9 @@ mod tests {
let cfg = VJepaConfig::small_test();
let mut trainer = VJepaTrainer::new(cfg, 5);
let m0 = trainer.train_step(2);
for _ in 0..4 { trainer.train_step(2); }
for _ in 0..4 {
trainer.train_step(2);
}
let m4 = trainer.train_step(2);
assert!(m4.ema_tau >= m0.ema_tau, "tau must anneal upward");
}
@@ -851,8 +961,8 @@ mod tests {
#[test]
fn test_neuro_config_patches() {
let cfg = NeuroJepaConfig::default();
assert_eq!(cfg.num_time_segments(), 16); // 1024 / 64
assert_eq!(cfg.total_patches(), 64 * 16); // 64 channels × 16 segments
assert_eq!(cfg.num_time_segments(), 16); // 1024 / 64
assert_eq!(cfg.total_patches(), 64 * 16); // 64 channels × 16 segments
}
#[test]
@@ -876,12 +986,20 @@ mod tests {
#[test]
fn test_neuro_mask_covers_all_patches() {
let cfg = NeuroJepaConfig { num_channels: 4, segment_len: 16, total_samples: 64, ..Default::default() };
let cfg = NeuroJepaConfig {
num_channels: 4,
segment_len: 16,
total_samples: 64,
..Default::default()
};
let mut strategy = NeuroMaskStrategy::new();
let result = strategy.generate(&cfg, 0);
let total = cfg.total_patches();
let covered = result.context_indices.len() + result.target_indices.len();
assert_eq!(covered, total, "all patches must appear in context or target");
assert_eq!(
covered, total,
"all patches must appear in context or target"
);
}
#[test]
@@ -919,7 +1037,10 @@ mod tests {
*channel_times.entry(channel).or_insert(0) += 1;
}
for (&_ch, &count) in &channel_times {
assert_eq!(count, ts, "tube masking: a masked channel must have ALL time segments masked");
assert_eq!(
count, ts,
"tube masking: a masked channel must have ALL time segments masked"
);
}
}
@@ -929,6 +1050,9 @@ mod tests {
let mut strategy = NeuroMaskStrategy::new();
let r0 = strategy.generate(&cfg, 0);
let r1 = strategy.generate(&cfg, 999);
assert_ne!(r0.target_indices, r1.target_indices, "different seeds should produce different masks");
assert_ne!(
r0.target_indices, r1.target_indices,
"different seeds should produce different masks"
);
}
}