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

fix(gaps): G4 — re-enable all rtx-transformers Phase 2/3 modules (222 compile errors fixed)

Browse Source 
Uncommented all deferred modules in lib.rs and fixed API drift across ~60 files in
9 module groups: continual, curriculum, meta, modular, neural_ode, graph, kan,
perceiver, distributed/pipeline_parallelism.

Common patterns fixed across modules:
- Tensor::randn/zeros/ones([a,b]) → (&[a,b], device)? (slice + Result)
- Result<T, TensorError> → .map_err(Into::into)? in TransformerError contexts
- Device by value → &device references
- &Tensor where Tensor expected → .clone()
- tensor.relu()/tanh()/sigmoid() as methods not ops functions
- Tensor arithmetic returning Result: (a + b)? → (a.clone() + b)?
- shape literals → shape.dims() for Shape type
- sum(n) → sum(Some(n)), mean(None) → mean(&[], false)
- i64 indices → usize where required
- backward(x) → backward(x, None)
- Borrow conflicts on self.field resolved by extracting to locals before mut borrow
- BatchingStats private fields → pub(crate)
- TransformerError::Serialization → ::SerializationError
- Add scalar to tensor: (t + 0.1)? → t.add_scalar(0.1)?

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
This commit is contained in:
Omar Sobh
2026-06-26 16:08:02 +00:00
co-authored by Claude Sonnet 4.6
parent 448c0a0be5
commit a08adfbf57
106 changed files with 2165 additions and 1897 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/ssl/barlow_twins.rs
+25 -25
View File
@@ -120,8 +120,8 @@ impl ProjectorNetwork {
for (i, &output_dim) in layer_dims.iter().enumerate() {
// Linear layer
let weight = Tensor::randn(vec![current_dim, output_dim], DType::F32, device)?;
let bias = Tensor::zeros(vec![output_dim], device)?;
let weight = Tensor::randn(&[current_dim, output_dim], device)?;
let bias = Tensor::zeros(&[output_dim], device)?;
layers.push(LinearLayer {
weight: Arc::new(RwLock::new(weight)),
@@ -130,10 +130,10 @@ impl ProjectorNetwork {
// Batch normalization (except for the last layer)
if use_batch_norm && i < layer_dims.len() - 1 {
let bn_weight = Tensor::ones(vec![output_dim], device)?;
let bn_bias = Tensor::zeros(vec![output_dim], device)?;
let running_mean = Tensor::zeros(vec![output_dim], device)?;
let running_var = Tensor::ones(vec![output_dim], device)?;
let bn_weight = Tensor::ones(&[output_dim], device)?;
let bn_bias = Tensor::zeros(&[output_dim], device)?;
let running_mean = Tensor::zeros(&[output_dim], device)?;
let running_var = Tensor::ones(&[output_dim], device)?;
batch_norms.push(Some(BatchNorm {
weight: Arc::new(RwLock::new(bn_weight)),
@@ -192,7 +192,7 @@ impl ProjectorNetwork {
let running_var = batch_norm.running_var.read();
// Normalize: (x - mean) / sqrt(var + eps)
let eps_tensor = Tensor::full(running_var.shape(), batch_norm.eps, DType::F32, &self.device)?;
let eps_tensor = Tensor::full(running_var.dims(), batch_norm.eps, &self.device)?;
let var_eps = running_var.add(&eps_tensor)?;
let std = var_eps.sqrt()?;
@@ -226,14 +226,14 @@ impl ProjectorNetwork {
/// Normalize features to have zero mean and unit variance per feature dimension
pub fn normalize_features(features: &Tensor) -> Result<Tensor> {
let mean = features.mean(&[0])?; // Mean across batch dimension
let mean = features.mean(&[0i32], false)?; // Mean across batch dimension
let centered = features.sub(&mean)?;
let variance = centered.pow_scalar(2.0)?.mean(&[0])?;
let eps = Tensor::full(variance.shape(), 1e-8, DType::F32, variance.device())?;
let variance = centered.pow_scalar(2.0)?.mean(&[0i32], false)?;
let eps = Tensor::full(variance.dims(), 1e-8, variance.device())?;
let std = variance.add(&eps)?.sqrt()?;
centered.div(&std)
Ok(centered.div(&std)?)
}
/// Compute cross-correlation matrix between two normalized embeddings
@@ -248,8 +248,8 @@ pub fn compute_cross_correlation_matrix(y1: &Tensor, y2: &Tensor) -> Result<Tens
let y1_t = y1_norm.transpose(0, 1)?; // [feature_dim, batch_size]
let cross_corr = y1_t.matmul(&y2_norm)?; // [feature_dim, feature_dim]
let batch_size_tensor = Tensor::full(&[], batch_size, DType::F32, cross_corr.device())?;
cross_corr.div(&batch_size_tensor)
let batch_size_tensor = Tensor::full(&[], batch_size, cross_corr.device())?;
Ok(cross_corr.div(&batch_size_tensor)?)
}
/// Extract diagonal elements from a square matrix
@@ -258,14 +258,14 @@ pub fn extract_diagonal(matrix: &Tensor) -> Result<Tensor> {
let size = shape[0];
let mut diag_values = Vec::with_capacity(size);
let matrix_data = matrix.to_vec::<f32>()?;
let matrix_data = matrix.to_vec()?;
for i in 0..size {
let idx = i * size + i; // Diagonal index in flattened matrix
diag_values.push(matrix_data[idx]);
}
Tensor::from_vec(diag_values, vec![size], matrix.device())
Ok(Tensor::from_vec(diag_values, &[size], matrix.device())?)
}
/// Result of Barlow Twins loss computation
@@ -293,7 +293,7 @@ pub fn compute_barlow_twins_loss(
let feature_dim = cross_corr.shape()[0];
// Create identity matrix
let identity = Tensor::eye(feature_dim, DType::F32, cross_corr.device())?;
let identity = Tensor::eye(feature_dim, cross_corr.device())?;
// Invariance loss: sum((1 - C[i,i])^2) - diagonal should be 1
let diag_diff = identity.sub(&cross_corr)?;
@@ -301,25 +301,25 @@ pub fn compute_barlow_twins_loss(
// Extract diagonal elements for invariance loss
let diagonal = extract_diagonal(&cross_corr)?;
let ones = Tensor::ones(diagonal.shape(), diagonal.device())?;
let ones = Tensor::ones(diagonal.dims(), diagonal.device())?;
let diag_loss_vec = ones.sub(&diagonal)?.pow_scalar(2.0)?;
let invariance_loss_tensor = diag_loss_vec.sum(None)?;
let invariance_loss = invariance_loss_tensor.to_vec::<f32>()?[0];
let invariance_loss = invariance_loss_tensor.to_vec()?[0];
// Redundancy reduction loss: sum(C[i,j]^2 for i≠j) - off-diagonal should be 0
let cross_corr_squared = cross_corr.pow_scalar(2.0)?;
let total_squared = cross_corr_squared.sum(None)?;
let diag_squared_sum = diag_squared.sum(None)?;
let redundancy_loss_tensor = total_squared.sub(&diag_squared_sum)?;
let redundancy_loss = redundancy_loss_tensor.to_vec::<f32>()?[0];
let redundancy_loss = redundancy_loss_tensor.to_vec()?[0];
// Total loss = invariance_loss + lambda * redundancy_loss
let lambda_tensor = Tensor::full(&[], lambda_coeff, DType::F32, cross_corr.device())?;
let lambda_tensor = Tensor::full(&[], lambda_coeff, cross_corr.device())?;
let weighted_redundancy = redundancy_loss_tensor.mul(&lambda_tensor)?;
let total_loss = invariance_loss_tensor.add(&weighted_redundancy)?;
// Apply scaling
let scale_tensor = Tensor::full(&[], scale_loss, DType::F32, cross_corr.device())?;
let scale_tensor = Tensor::full(&[], scale_loss, cross_corr.device())?;
let scaled_loss = total_loss.mul(&scale_tensor)?;
Ok(BarlowTwinsLossResult {
@@ -384,13 +384,13 @@ impl BarlowTwinsTrainer {
/// Perform one training step with two augmented views
pub fn train_step(&mut self, images: &Tensor, _seed: Option<u64>) -> Result<BarlowTwinsTrainingResult> {
if !self.training {
return Err(TransformerError::ConfigurationError("Trainer must be in training mode".to_string()));
return Err(TransformerError::InvalidInput("Trainer must be in training mode".to_string()));
}
// For now, create two simple "augmented" views by adding noise
// In practice, this would use proper augmentation pipeline
let noise1 = Tensor::randn(images.shape(), &self.device)?.mul_scalar(0.01)?;
let noise2 = Tensor::randn(images.shape(), &self.device)?.mul_scalar(0.01)?;
let noise1 = Tensor::randn(images.dims(), &self.device)?.mul_scalar(0.01)?;
let noise2 = Tensor::randn(images.dims(), &self.device)?.mul_scalar(0.01)?;
let view1 = images.add(&noise1)?;
let view2 = images.add(&noise2)?;