rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 20 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/simmim_simplified.rs
+146 -58
View File
@@ -4,8 +4,8 @@
//! It focuses on the algorithmic structure and can be upgraded when tensor operations are stable.
use crate::prelude::*;
use std::sync::Arc;
use parking_lot::RwLock;
use std::sync::Arc;
/// Mock tensor for simplified implementation
#[derive(Debug, Clone)]
@@ -24,7 +24,7 @@ impl SimpleTensor {
device: device.clone(),
})
}
pub fn randn(shape: Vec<usize>, device: &Device) -> Result<Self> {
let size = shape.iter().product();
let mut data = Vec::with_capacity(size);
@@ -34,30 +34,46 @@ impl SimpleTensor {
let val = (seed as f32 / u32::MAX as f32) * 2.0 - 1.0;
data.push(val);
}
Ok(Self { data, shape, device: device.clone() })
Ok(Self {
data,
shape,
device: device.clone(),
})
}
pub fn from_data(data: Vec<f32>, shape: Vec<usize>, device: &Device) -> Result<Self> {
Ok(Self { data, shape, device: device.clone() })
Ok(Self {
data,
shape,
device: device.clone(),
})
}
pub fn shape(&self) -> &[usize] { &self.shape }
pub fn device(&self) -> &Device { &self.device }
pub fn to_vec<T>(&self) -> Result<Vec<T>>
where
T: Clone + From<f32>
pub fn shape(&self) -> &[usize] {
&self.shape
}
pub fn device(&self) -> &Device {
&self.device
}
pub fn to_vec<T>(&self) -> Result<Vec<T>>
where
T: Clone + From<f32>,
{
Ok(self.data.iter().map(|&x| T::from(x)).collect())
}
pub fn matmul(&self, _other: &Self) -> Result<Self> {
// Simplified matrix multiplication - just return a tensor of appropriate shape
let result_shape = vec![self.shape[0], _other.shape[1]];
let result_data = vec![1.0; result_shape.iter().product()];
Ok(Self { data: result_data, shape: result_shape, device: self.device.clone() })
Ok(Self {
data: result_data,
shape: result_shape,
device: self.device.clone(),
})
}
pub fn add(&self, other: &Self) -> Result<Self> {
let mut result_data = self.data.clone();
for (i, &val) in other.data.iter().enumerate() {
@@ -65,9 +81,13 @@ impl SimpleTensor {
result_data[i] += val;
}
}
Ok(Self { data: result_data, shape: self.shape.clone(), device: self.device.clone() })
Ok(Self {
data: result_data,
shape: self.shape.clone(),
device: self.device.clone(),
})
}
pub fn sub(&self, other: &Self) -> Result<Self> {
let mut result_data = self.data.clone();
for (i, &val) in other.data.iter().enumerate() {
@@ -75,28 +95,52 @@ impl SimpleTensor {
result_data[i] -= val;
}
}
Ok(Self { data: result_data, shape: self.shape.clone(), device: self.device.clone() })
Ok(Self {
data: result_data,
shape: self.shape.clone(),
device: self.device.clone(),
})
}
pub fn abs(&self) -> Result<Self> {
let result_data = self.data.iter().map(|x| x.abs()).collect();
Ok(Self { data: result_data, shape: self.shape.clone(), device: self.device.clone() })
Ok(Self {
data: result_data,
shape: self.shape.clone(),
device: self.device.clone(),
})
}
pub fn pow_scalar(&self, exp: f32) -> Result<Self> {
let result_data = self.data.iter().map(|x| x.powf(exp)).collect();
Ok(Self { data: result_data, shape: self.shape.clone(), device: self.device.clone() })
Ok(Self {
data: result_data,
shape: self.shape.clone(),
device: self.device.clone(),
})
}
pub fn mean(&self, _dims: &[usize]) -> Result<Self> {
let mean_val = if self.data.is_empty() { 0.0 } else { self.data.iter().sum::<f32>() / self.data.len() as f32 };
Ok(Self { data: vec![mean_val], shape: vec![], device: self.device.clone() })
let mean_val = if self.data.is_empty() {
0.0
} else {
self.data.iter().sum::<f32>() / self.data.len() as f32
};
Ok(Self {
data: vec![mean_val],
shape: vec![],
device: self.device.clone(),
})
}
pub fn reshape(&self, new_shape: &[usize]) -> Result<Self> {
Ok(Self { data: self.data.clone(), shape: new_shape.to_vec(), device: self.device.clone() })
Ok(Self {
data: self.data.clone(),
shape: new_shape.to_vec(),
device: self.device.clone(),
})
}
pub fn full(shape: &[usize], value: f32, device: &Device) -> Result<Self> {
let size = shape.iter().product();
Ok(Self {
@@ -245,17 +289,25 @@ impl RandomPatchMasker {
}
/// Generate random patch mask for a batch
pub fn mask_patches(&self, patches: &SimpleTensor, seed: Option<u64>) -> Result<PatchMaskResult> {
pub fn mask_patches(
&self,
patches: &SimpleTensor,
seed: Option<u64>,
) -> Result<PatchMaskResult> {
let shape = patches.shape();
if shape.len() != 3 {
return Err(TransformerError::InvalidInput("Patches tensor must be 3D [batch, patches, embed_dim]".into()));
return Err(TransformerError::InvalidInput(
"Patches tensor must be 3D [batch, patches, embed_dim]".into(),
));
}
let batch_size = shape[0];
let num_patches = shape[1];
if num_patches == 0 {
return Err(TransformerError::InvalidInput("Number of patches cannot be zero".into()));
return Err(TransformerError::InvalidInput(
"Number of patches cannot be zero".into(),
));
}
let num_masked = (num_patches as f32 * self.mask_ratio).round() as usize;
@@ -263,13 +315,13 @@ impl RandomPatchMasker {
let mut mask_indices = Vec::new();
let mut visible_indices = Vec::new();
let mut rng_state = seed.unwrap_or(42);
for _ in 0..batch_size {
// Create all patch indices
let mut indices: Vec<usize> = (0..num_patches).collect();
// Shuffle indices deterministically
for i in (1..indices.len()).rev() {
rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
@@ -401,14 +453,18 @@ pub struct SimMIMLoss;
impl SimMIMLoss {
/// Compute L1 (MAE) loss between predicted and target pixels
pub fn compute_l1_loss(predicted: &SimpleTensor, target: &SimpleTensor) -> Result<SimpleTensor> {
pub fn compute_l1_loss(
predicted: &SimpleTensor,
target: &SimpleTensor,
) -> Result<SimpleTensor> {
let pred_shape = predicted.shape();
let target_shape = target.shape();
if pred_shape != target_shape {
return Err(TransformerError::InvalidInput(
format!("Predicted and target shapes must match: {:?} vs {:?}", pred_shape, target_shape)
));
return Err(TransformerError::InvalidInput(format!(
"Predicted and target shapes must match: {:?} vs {:?}",
pred_shape, target_shape
)));
}
let diff = predicted.sub(target)?;
@@ -417,19 +473,26 @@ impl SimMIMLoss {
}
/// Compute normalized L1 loss
pub fn compute_l1_loss_normalized(predicted: &SimpleTensor, target: &SimpleTensor) -> Result<SimpleTensor> {
pub fn compute_l1_loss_normalized(
predicted: &SimpleTensor,
target: &SimpleTensor,
) -> Result<SimpleTensor> {
Self::compute_l1_loss(predicted, target)
}
/// Compute L2 (MSE) loss between predicted and target pixels
pub fn compute_l2_loss(predicted: &SimpleTensor, target: &SimpleTensor) -> Result<SimpleTensor> {
pub fn compute_l2_loss(
predicted: &SimpleTensor,
target: &SimpleTensor,
) -> Result<SimpleTensor> {
let pred_shape = predicted.shape();
let target_shape = target.shape();
if pred_shape != target_shape {
return Err(TransformerError::InvalidInput(
format!("Predicted and target shapes must match: {:?} vs {:?}", pred_shape, target_shape)
));
return Err(TransformerError::InvalidInput(format!(
"Predicted and target shapes must match: {:?} vs {:?}",
pred_shape, target_shape
)));
}
let diff = predicted.sub(target)?;
@@ -471,10 +534,10 @@ impl SimMIMTrainer {
) -> Result<Self> {
let masker = RandomPatchMasker::new(config.mask_ratio);
let mask_token = MaskTokenEmbedding::new(config.encoder_dim, device)?;
let patch_volume = config.patch_size * config.patch_size * in_channels;
let prediction_head = LinearPredictionHead::new(config.encoder_dim, patch_volume, device)?;
let encoder = SimpleTensor::randn(vec![patch_volume, config.encoder_dim], device)?;
Ok(Self {
@@ -514,7 +577,11 @@ impl SimMIMTrainer {
}
/// Perform one SimMIM training step
pub fn train_step(&mut self, images: &SimpleTensor, seed: Option<u64>) -> Result<SimMIMTrainingResult> {
pub fn train_step(
&mut self,
images: &SimpleTensor,
seed: Option<u64>,
) -> Result<SimMIMTrainingResult> {
let patches = self.extract_patches(images)?;
let target_pixels = patches.clone();
@@ -528,14 +595,15 @@ impl SimMIMTrainer {
let loss = match self.config.loss_type {
LossType::L1 => {
if self.config.norm_pix_loss {
SimMIMLoss::compute_l1_loss_normalized(&predicted_pixels, &target_masked_pixels)?
SimMIMLoss::compute_l1_loss_normalized(
&predicted_pixels,
&target_masked_pixels,
)?
} else {
SimMIMLoss::compute_l1_loss(&predicted_pixels, &target_masked_pixels)?
}
}
LossType::L2 => {
SimMIMLoss::compute_l2_loss(&predicted_pixels, &target_masked_pixels)?
}
LossType::L2 => SimMIMLoss::compute_l2_loss(&predicted_pixels, &target_masked_pixels)?,
};
Ok(SimMIMTrainingResult {
@@ -568,7 +636,11 @@ impl SimMIMTrainer {
images.reshape(&[batch_size, num_patches, patch_volume])
}
fn apply_mask_tokens(&self, patches: &SimpleTensor, mask_result: &PatchMaskResult) -> Result<SimpleTensor> {
fn apply_mask_tokens(
&self,
patches: &SimpleTensor,
mask_result: &PatchMaskResult,
) -> Result<SimpleTensor> {
// Simplified implementation - just return the original patches
Ok(patches.clone())
}
@@ -578,23 +650,39 @@ impl SimMIMTrainer {
patches.matmul(&*encoder)
}
fn extract_masked_features(&self, features: &SimpleTensor, mask_result: &PatchMaskResult) -> Result<SimpleTensor> {
fn extract_masked_features(
&self,
features: &SimpleTensor,
mask_result: &PatchMaskResult,
) -> Result<SimpleTensor> {
let shape = features.shape();
let batch_size = shape[0];
let embed_dim = shape[2];
let num_masked = mask_result.num_masked_patches;
let result_data = vec![1.0f32; batch_size * num_masked * embed_dim];
SimpleTensor::from_data(result_data, vec![batch_size, num_masked, embed_dim], &self.device)
SimpleTensor::from_data(
result_data,
vec![batch_size, num_masked, embed_dim],
&self.device,
)
}
fn extract_masked_targets(&self, targets: &SimpleTensor, mask_result: &PatchMaskResult) -> Result<SimpleTensor> {
fn extract_masked_targets(
&self,
targets: &SimpleTensor,
mask_result: &PatchMaskResult,
) -> Result<SimpleTensor> {
let shape = targets.shape();
let batch_size = shape[0];
let patch_volume = shape[2];
let num_masked = mask_result.num_masked_patches;
let result_data = vec![0.5f32; batch_size * num_masked * patch_volume];
SimpleTensor::from_data(result_data, vec![batch_size, num_masked, patch_volume], &self.device)
SimpleTensor::from_data(
result_data,
vec![batch_size, num_masked, patch_volume],
&self.device,
)
}
}
}