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

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]>
2026-08-10 07:09:36 -07:00
parent ad6405663f
commit 4aaa36a57a
305 changed files with 25537 additions and 18337 deletions
@@ -2,7 +2,7 @@
 //!
 //! Comprehensive tests for VICReg (Variance-Invariance-Covariance Regularization) implementation.
 //! Follows strict TDD - tests are written first, then implementation.
-//! 
+//!
 //! VICReg is a self-supervised learning method that explicitly avoids collapse through:
 //! - Invariance: Similar representations for augmented views (MSE loss)  
 //! - Variance: Maintains variance ≥ γ in each dimension (hinge loss)
@@ -10,21 +10,21 @@

 #[cfg(all(test, feature = "disabled_tests"))]
 mod tests {
-    use crate::prelude::*;
-    use super::super::vicreg::*;
    use super::super::byol::Backbone;
+    use super::super::vicreg::*;
+    use crate::prelude::*;

    /// Test VICRegConfig creation with default hyperparameters
    #[test]
    fn test_vicreg_config_default() {
        let config = VICRegConfig::default();
-        
+
        // Default hyperparameters from VICReg paper
-        assert_eq!(config.sim_coeff, 25.0);        // λ - invariance weight
-        assert_eq!(config.std_coeff, 25.0);        // μ - variance weight  
-        assert_eq!(config.cov_coeff, 1.0);         // ν - covariance weight
-        assert_eq!(config.variance_target, 1.0);   // γ - target standard deviation
-        assert_eq!(config.epsilon, 1e-4);          // ε - numerical stability
+        assert_eq!(config.sim_coeff, 25.0); // λ - invariance weight
+        assert_eq!(config.std_coeff, 25.0); // μ - variance weight  
+        assert_eq!(config.cov_coeff, 1.0); // ν - covariance weight
+        assert_eq!(config.variance_target, 1.0); // γ - target standard deviation
+        assert_eq!(config.epsilon, 1e-4); // ε - numerical stability
        assert_eq!(config.expander_dims, vec![8192, 8192, 8192]);
    }

@@ -37,7 +37,7 @@ mod tests {
            .with_cov_coeff(2.0)
            .with_variance_target(0.5)
            .with_epsilon(1e-5);
-        
+
        assert_eq!(config.backbone_dim, 2048);
        assert_eq!(config.expander_dims, vec![4096, 4096, 2048]);
        assert_eq!(config.sim_coeff, 20.0);
@@ -45,7 +45,7 @@ mod tests {
        assert_eq!(config.cov_coeff, 2.0);
        assert_eq!(config.variance_target, 0.5);
        assert_eq!(config.epsilon, 1e-5);
-        
+
        Ok(())
    }

@@ -55,13 +55,13 @@ mod tests {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let input_dim = 2048;
        let dims = vec![4096, 4096, 2048];
-        
+
        let expander = ExpanderNetwork::new(input_dim, dims.clone(), &device)?;
-        
+
        assert_eq!(expander.input_dim(), input_dim);
        assert_eq!(expander.output_dim(), *dims.last().unwrap());
        assert_eq!(expander.layer_dims(), &dims);
-        
+
        Ok(())
    }

@@ -73,14 +73,14 @@ mod tests {
        let input_dim = 2048;
        let output_dim = 8192;
        let dims = vec![4096, 4096, output_dim];
-        
+
        let expander = ExpanderNetwork::new(input_dim, dims, &device)?;
        let input = Tensor::randn(vec![batch_size, input_dim], DType::F32, &device)?;
-        
+
        let output = expander.forward(&input)?;
-        
+
        assert_eq!(output.shape(), &[batch_size, output_dim]);
-        
+
        Ok(())
    }

@@ -90,20 +90,20 @@ mod tests {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 64;
        let feature_dim = 8192;
-        
+
        // Create two representations
        let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
        let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
-        
+
        let inv_loss = compute_invariance_loss(&y1, &y2)?;
-        
+
        // Invariance loss should be a scalar (MSE)
        assert_eq!(inv_loss.shape(), &[]);
-        
-        // Loss should be positive 
+
+        // Loss should be positive
        let loss_val = inv_loss.to_scalar::<f32>()?;
        assert!(loss_val >= 0.0);
-        
+
        Ok(())
    }

@@ -115,19 +115,19 @@ mod tests {
        let feature_dim = 8192;
        let gamma = 1.0;
        let epsilon = 1e-4;
-        
+
        // Create representation
        let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
-        
+
        let var_loss = compute_variance_loss(&y, gamma, epsilon)?;
-        
+
        // Variance loss should be a scalar
        assert_eq!(var_loss.shape(), &[]);
-        
+
        // Loss should be non-negative (hinge loss)
        let loss_val = var_loss.to_scalar::<f32>()?;
        assert!(loss_val >= 0.0);
-        
+
        Ok(())
    }

@@ -138,19 +138,19 @@ mod tests {
        let batch_size = 64;
        let feature_dim = 128; // Smaller for computational efficiency in tests
        let epsilon = 1e-4;
-        
+
        // Create representation
        let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
-        
+
        let cov_loss = compute_covariance_loss(&y, epsilon)?;
-        
+
        // Covariance loss should be a scalar
        assert_eq!(cov_loss.shape(), &[]);
-        
+
        // Loss should be non-negative (sum of squared off-diagonal elements)
        let loss_val = cov_loss.to_scalar::<f32>()?;
        assert!(loss_val >= 0.0);
-        
+
        Ok(())
    }

@@ -160,15 +160,15 @@ mod tests {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 64;
        let feature_dim = 8192;
-        
+
        let config = VICRegConfig::default();
-        
+
        // Create two representations (from augmented views)
        let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
        let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
-        
+
        let loss_result = compute_vicreg_loss(&y1, &y2, &config)?;
-        
+
        // Check that all loss components are computed
        assert_eq!(loss_result.total_loss.shape(), &[]);
        assert_eq!(loss_result.invariance_loss.shape(), &[]);
@@ -176,7 +176,7 @@ mod tests {
        assert_eq!(loss_result.variance_loss_y2.shape(), &[]);
        assert_eq!(loss_result.covariance_loss_y1.shape(), &[]);
        assert_eq!(loss_result.covariance_loss_y2.shape(), &[]);
-        
+
        // All losses should be non-negative
        let total = loss_result.total_loss.to_scalar::<f32>()?;
        let inv = loss_result.invariance_loss.to_scalar::<f32>()?;
@@ -184,21 +184,21 @@ mod tests {
        let var2 = loss_result.variance_loss_y2.to_scalar::<f32>()?;
        let cov1 = loss_result.covariance_loss_y1.to_scalar::<f32>()?;
        let cov2 = loss_result.covariance_loss_y2.to_scalar::<f32>()?;
-        
+
        assert!(total >= 0.0);
        assert!(inv >= 0.0);
        assert!(var1 >= 0.0);
        assert!(var2 >= 0.0);
        assert!(cov1 >= 0.0);
        assert!(cov2 >= 0.0);
-        
+
        // Total loss should be weighted sum of components
-        let expected_total = config.sim_coeff * inv + 
-                           config.std_coeff * (var1 + var2) +
-                           config.cov_coeff * (cov1 + cov2);
-        
+        let expected_total = config.sim_coeff * inv
+            + config.std_coeff * (var1 + var2)
+            + config.cov_coeff * (cov1 + cov2);
+
        assert!((total - expected_total).abs() < 1e-3);
-        
+
        Ok(())
    }

@@ -207,34 +207,38 @@ mod tests {
    fn test_vicreg_trainer_init() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = VICRegConfig::default();
-        
-        // Create simple backbone mock 
+
+        // Create simple backbone mock
        struct SimpleBackbone {
            output_dim: usize,
            device: Device,
        }
-        
+
        impl Backbone for SimpleBackbone {
            fn forward(&self, x: &Tensor) -> Result<Tensor> {
                let batch_size = x.shape()[0];
-                Ok(Tensor::randn(vec![batch_size, self.output_dim], DType::F32, &self.device)?)
+                Ok(Tensor::randn(
+                    vec![batch_size, self.output_dim],
+                    DType::F32,
+                    &self.device,
+                )?)
            }
-            
+
            fn output_dim(&self) -> usize {
                self.output_dim
            }
        }
-        
+
        let backbone = SimpleBackbone {
            output_dim: 2048,
            device: device.clone(),
        };
-        
+
        let trainer = VICRegTrainer::new(backbone, config, &device)?;
-        
+
        assert_eq!(trainer.config().backbone_dim, 2048);
        assert_eq!(trainer.config().expander_dims, vec![8192, 8192, 8192]);
-        
+
        Ok(())
    }

@@ -244,85 +248,93 @@ mod tests {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 32;
        let config = VICRegConfig::default();
-        
+
        struct SimpleBackbone {
            output_dim: usize,
            device: Device,
        }
-        
+
        impl Backbone for SimpleBackbone {
            fn forward(&self, x: &Tensor) -> Result<Tensor> {
                let batch_size = x.shape()[0];
-                Ok(Tensor::randn(vec![batch_size, self.output_dim], DType::F32, &self.device)?)
+                Ok(Tensor::randn(
+                    vec![batch_size, self.output_dim],
+                    DType::F32,
+                    &self.device,
+                )?)
            }
-            
+
            fn output_dim(&self) -> usize {
                self.output_dim
            }
        }
-        
+
        let backbone = SimpleBackbone {
            output_dim: 2048,
            device: device.clone(),
        };
-        
+
        let trainer = VICRegTrainer::new(backbone, config.clone(), &device)?;
-        
+
        // Create two augmented views
        let x1 = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
        let x2 = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
-        
+
        let result = trainer.forward(&x1, &x2)?;
-        
+
        // Should return VICReg loss components
        assert_eq!(result.total_loss.shape(), &[]);
        assert!(result.total_loss.to_scalar::<f32>()? >= 0.0);
-        
+
        Ok(())
    }

-    /// Test training step with gradient updates 
+    /// Test training step with gradient updates
    #[test]
    fn test_vicreg_training_step() -> Result<()> {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 16; // Small batch for test
        let config = VICRegConfig::default();
-        
+
        struct SimpleBackbone {
            output_dim: usize,
            device: Device,
        }
-        
+
        impl Backbone for SimpleBackbone {
            fn forward(&self, x: &Tensor) -> Result<Tensor> {
                let batch_size = x.shape()[0];
-                Ok(Tensor::randn(vec![batch_size, self.output_dim], DType::F32, &self.device)?)
+                Ok(Tensor::randn(
+                    vec![batch_size, self.output_dim],
+                    DType::F32,
+                    &self.device,
+                )?)
            }
-            
+
            fn output_dim(&self) -> usize {
                self.output_dim
            }
        }
-        
+
        let backbone = SimpleBackbone {
            output_dim: 2048,
            device: device.clone(),
        };
-        
+
        let mut trainer = VICRegTrainer::new(backbone, config, &device)?;
-        
+
        // Create batch of images (normally these would be augmented views)
        let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
-        
+
        // Perform training step
        let result = trainer.train_step(&images)?;
-        
+
        // Should return training metrics
        assert!(result.total_loss >= 0.0);
        assert!(result.invariance_loss >= 0.0);
        assert!(result.variance_loss >= 0.0);
        assert!(result.covariance_loss >= 0.0);
-        
+
        Ok(())
    }

@@ -332,16 +344,16 @@ mod tests {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 32;
        let feature_dim = 256;
-        
+
        // Identical representations
        let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
-        
+
        let inv_loss = compute_invariance_loss(&y, &y)?;
        let loss_val = inv_loss.to_scalar::<f32>()?;
-        
+
        // Invariance loss should be very close to zero for identical inputs
        assert!(loss_val < 1e-6);
-        
+
        Ok(())
    }

@@ -353,18 +365,18 @@ mod tests {
        let feature_dim = 256;
        let gamma = 1.0;
        let epsilon = 1e-4;
-        
+
        // Create high variance input (scale up random values)
        let y_base = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
        let scale = Tensor::full(y_base.shape(), 3.0, DType::F32, &device)?;
        let y = y_base.mul(&scale)?;
-        
+
        let var_loss = compute_variance_loss(&y, gamma, epsilon)?;
        let loss_val = var_loss.to_scalar::<f32>()?;
-        
+
        // High variance should result in low variance loss (hinge = max(0, γ - std))
        assert!(loss_val < 0.1);
-        
+
        Ok(())
    }

@@ -375,25 +387,25 @@ mod tests {
        let batch_size = 100;
        let feature_dim = 10; // Small for controlled test
        let epsilon = 1e-4;
-        
+
        // Create orthogonal/uncorrelated features manually
        let mut y_data = vec![0.0; batch_size * feature_dim];
-        
+
        // Fill each feature dimension independently
        for i in 0..feature_dim {
            for b in 0..batch_size {
                y_data[b * feature_dim + i] = (b as f32 + i as f32 * 100.0).sin();
            }
        }
-        
+
        let y = Tensor::from_slice(&y_data, vec![batch_size, feature_dim], DType::F32, &device)?;
-        
+
        let cov_loss = compute_covariance_loss(&y, epsilon)?;
        let loss_val = cov_loss.to_scalar::<f32>()?;
-        
+
        // Uncorrelated features should have relatively low covariance loss
        assert!(loss_val >= 0.0); // Should be non-negative
-        
+
        Ok(())
    }

@@ -404,26 +416,26 @@ mod tests {
        let feature_dim = 512;
        let epsilon = 1e-4;
        let gamma = 1.0;
-        
+
        // Test different batch sizes
        for batch_size in [1, 8, 32, 128] {
            let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
            let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
-            
+
            // All loss computations should work regardless of batch size
            let inv_loss = compute_invariance_loss(&y1, &y2)?;
            let var_loss = compute_variance_loss(&y1, gamma, epsilon)?;
            let cov_loss = compute_covariance_loss(&y1, epsilon)?;
-            
+
            assert_eq!(inv_loss.shape(), &[]);
            assert_eq!(var_loss.shape(), &[]);
            assert_eq!(cov_loss.shape(), &[]);
-            
+
            assert!(inv_loss.to_scalar::<f32>()? >= 0.0);
            assert!(var_loss.to_scalar::<f32>()? >= 0.0);
            assert!(cov_loss.to_scalar::<f32>()? >= 0.0);
        }
-        
+
        Ok(())
    }
-}
+}