rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 29 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/mean_teacher.rs
+29 -29
View File
@@ -129,9 +129,9 @@ impl LinearLayer {
fn new(input_dim: usize, output_dim: usize, device: &Device) -> Result<Self> {
// Xavier initialization
let scale = (2.0 / (input_dim + output_dim) as f32).sqrt();
let weight = Tensor::randn(vec![input_dim, output_dim], DType::F32, device)?
let weight = Tensor::randn(&[input_dim, output_dim], device)?
.mul_scalar(scale)?;
let bias = Tensor::zeros(vec![output_dim], device)?;
let bias = Tensor::zeros(&[output_dim], device)?;
Ok(Self {
weight: Arc::new(RwLock::new(weight)),
@@ -144,7 +144,7 @@ impl LinearLayer {
let bias = self.bias.read();
let output = input.matmul(&*weight)?;
output.add(&*bias)
Ok(output.add(&*bias)?)
}
fn get_weight(&self) -> Tensor {
@@ -388,20 +388,21 @@ impl NoiseAugmenter {
}
fn apply_gaussian_noise(&self, input: &Tensor, _seed: Option<u64>) -> Result<Tensor> {
let noise = Tensor::randn(input.shape().clone(), &self.device)?
let noise = Tensor::randn(input.dims(), &self.device)?
.mul_scalar(self.noise_level)?;
input.add(&noise)
Ok(input.add(&noise)?)
}
fn apply_dropout_noise(&self, input: &Tensor, _seed: Option<u64>) -> Result<Tensor> {
// Simulate dropout by randomly scaling elements
let keep_prob = 1.0 - self.noise_level;
let mask = Tensor::rand(input.shape().clone(), DType::F32, &self.device)?;
// Use randn sigmoid to get uniform [0,1]-like mask
let mask = Tensor::randn(input.dims(), &self.device)?.sigmoid()?;
let dropout_mask = mask.gt_scalar(self.noise_level)?
.to_dtype(DType::F32)?
.div_scalar(keep_prob)?;
input.mul(&dropout_mask)
Ok(input.mul(&dropout_mask)?)
}
}
@@ -409,7 +410,10 @@ impl NoiseAugmenter {
pub fn compute_consistency_loss(student_pred: &Tensor, teacher_pred: &Tensor) -> Result<Tensor> {
let diff = student_pred.sub(teacher_pred)?;
let squared_diff = diff.mul(&diff)?;
squared_diff.mean(None, false)
// compute mean over all elements by summing and dividing
let n = squared_diff.dims().iter().product::<usize>() as f32;
let total = squared_diff.sum(None)?;
Ok(total.div_scalar(n)?)
}
/// Consistency weight ramp-up scheduler
@@ -574,9 +578,9 @@ impl MeanTeacherTrainer {
self.update_teacher_parameters()?;
Ok(MeanTeacherTrainingResult {
supervised_loss: supervised_loss.to_vec::<f32>()?[0],
supervised_loss: supervised_loss.to_vec()?[0],
consistency_loss: 0.0,
total_loss: supervised_loss.to_vec::<f32>()?[0],
total_loss: supervised_loss.to_vec()?[0],
})
}
@@ -608,8 +612,8 @@ impl MeanTeacherTrainer {
Ok(MeanTeacherTrainingResult {
supervised_loss: 0.0,
consistency_loss: weighted_consistency_loss.to_vec::<f32>()?[0],
total_loss: weighted_consistency_loss.to_vec::<f32>()?[0],
consistency_loss: weighted_consistency_loss.to_vec()?[0],
total_loss: weighted_consistency_loss.to_vec()?[0],
})
}
@@ -645,9 +649,9 @@ impl MeanTeacherTrainer {
self.update_teacher_parameters()?;
Ok(MeanTeacherTrainingResult {
supervised_loss: supervised_loss.to_vec::<f32>()?[0],
consistency_loss: weighted_consistency_loss.to_vec::<f32>()?[0],
total_loss: total_loss.to_vec::<f32>()?[0],
supervised_loss: supervised_loss.to_vec()?[0],
consistency_loss: weighted_consistency_loss.to_vec()?[0],
total_loss: total_loss.to_vec()?[0],
})
}
@@ -660,26 +664,22 @@ impl MeanTeacherTrainer {
// Simplified cross-entropy loss
let log_probs = predictions.log_softmax(-1)?;
let labels_one_hot = self.to_one_hot(labels, predictions.shape()[1])?;
let loss = log_probs.mul(&labels_one_hot)?.sum(None, false)?.neg()?;
loss.div_scalar(predictions.shape()[0] as f32)
let loss = log_probs.mul(&labels_one_hot)?.sum(None)?.neg()?;
Ok(loss.div_scalar(predictions.shape()[0] as f32)?)
}
fn to_one_hot(&self, labels: &Tensor, num_classes: usize) -> Result<Tensor> {
let batch_size = labels.shape()[0];
let mut one_hot = Tensor::zeros(vec![batch_size, num_classes], DType::F32, &self.device)?;
// Simplified one-hot encoding (would need proper indexing in real implementation)
for i in 0..batch_size {
let label_val = labels.get(i)?.to_vec::<i64>()?[0] as usize;
if label_val < num_classes {
one_hot = one_hot.index_put(
&[Some(i)],
&Tensor::ones(vec![1], &self.device)?
)?;
// Simplified one-hot: build the data manually then create tensor
let labels_data = labels.to_vec()?;
let mut one_hot_data = vec![0.0f32; batch_size * num_classes];
for (i, &label_val) in labels_data.iter().enumerate().take(batch_size) {
let idx = label_val as usize;
if idx < num_classes {
one_hot_data[i * num_classes + idx] = 1.0;
}
}
Ok(one_hot)
Ok(Tensor::from_vec(one_hot_data, &[batch_size, num_classes], &self.device)?)
}
fn update_teacher_parameters(&mut self) -> Result<()> {