rustytorch/crates/training/rtx-transformers/examples/normalization_demo.rs

//! Advanced Normalization Techniques Demo
//!
//! This example demonstrates the comprehensive normalization framework
//! implemented in RTX-Transformers, showcasing all available normalization
//! techniques and their use cases.

use rtx_tensor::{Device, Shape, Tensor};
use rtx_transformers::{
    Result,
    layers::{
        AdaLayerNorm,
        DeepNorm,
        DeepNormConfig,
        // Layer trait
        Layer,
        // Individual normalization types
        LayerNorm,
        // Position wrappers
        NormPosition,
        NormWrapper,

        NormalizationConfig,
        // Unified framework
        NormalizationType,
        PostNorm,
        PowerNorm,
        PreNorm,
        RMSNorm,
        ScaleNorm,
        UnifiedNorm,
        normalization_factory,

        power_norm_presets,
    },
};
use tracing::{Level, info, warn};

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt().with_max_level(Level::INFO).init();

    info!("🚀 RTX Advanced Normalization Techniques Demo");

    // Setup device
    let device = Device::cuda(0)?;
    let d_model = 512;
    let batch_size = 2;
    let seq_len = 8;

    // Create sample input tensor
    let input = Tensor::randn(Shape::new(vec![batch_size, seq_len, d_model])?, &device)?;

    info!("📊 Input tensor shape: {:?}", input.shape().dims());

    // Demo 1: Individual Normalization Techniques
    demo_individual_norms(&input, &device).await?;

    // Demo 2: Unified Normalization Framework
    demo_unified_framework(&input, &device).await?;

    // Demo 3: PreNorm vs PostNorm Configurations
    demo_norm_positions(&input, &device).await?;

    // Demo 4: DeepNorm for Very Deep Networks
    demo_deep_norm(&input, &device).await?;

    // Demo 5: Adaptive Normalization with Conditioning
    demo_adaptive_norm(&input, &device).await?;

    // Demo 6: Performance Comparison
    demo_performance_comparison(&input, &device).await?;

    info!("✅ Demo completed successfully!");
    Ok(())
}

/// Demonstrate individual normalization techniques
async fn demo_individual_norms(input: &Tensor, device: &Device) -> Result<()> {
    info!("\n🔬 Demo 1: Individual Normalization Techniques");

    let d_model = input.shape().dims()[2];

    // 1. Standard LayerNorm
    info!("  Testing LayerNorm...");
    let layer_norm = LayerNorm::new(d_model, 1e-5, true, device)?;
    let ln_output = layer_norm.forward(input)?;
    info!("    LayerNorm output shape: {:?}", ln_output.shape().dims());

    // 2. RMSNorm (used in modern LLMs)
    info!("  Testing RMSNorm...");
    let rms_norm = RMSNorm::new(d_model, 1e-6, device)?;
    let rms_output = rms_norm.forward(input)?;
    info!("    RMSNorm output shape: {:?}", rms_output.shape().dims());

    // 3. ScaleNorm (efficient alternative)
    info!("  Testing ScaleNorm...");
    let scale_norm = ScaleNorm::new(d_model, 1e-8, device)?;
    let scale_output = scale_norm.forward(input)?;
    info!(
        "    ScaleNorm output shape: {:?}",
        scale_output.shape().dims()
    );

    // 4. PowerNorm variants
    info!("  Testing PowerNorm variants...");

    // L1 norm (Manhattan)
    let l1_norm = power_norm_presets::manhattan_norm(d_model, device)?;
    let l1_output = l1_norm.forward(input)?;
    info!("    L1 norm output shape: {:?}", l1_output.shape().dims());

    // L2 norm (Euclidean, equivalent to RMSNorm)
    let l2_norm = power_norm_presets::euclidean_norm(d_model, device)?;
    let l2_output = l2_norm.forward(input)?;
    info!("    L2 norm output shape: {:?}", l2_output.shape().dims());

    // Max norm (L∞)
    let max_norm = power_norm_presets::max_norm(d_model, device)?;
    let max_output = max_norm.forward(input)?;
    info!("    Max norm output shape: {:?}", max_output.shape().dims());

    Ok(())
}

/// Demonstrate unified normalization framework
async fn demo_unified_framework(input: &Tensor, device: &Device) -> Result<()> {
    info!("\n🎯 Demo 2: Unified Normalization Framework");

    let d_model = input.shape().dims()[2];

    // Create different normalization configurations
    let configs = vec![
        (
            "Standard LayerNorm",
            normalization_factory::layer_norm(d_model),
        ),
        ("Modern RMSNorm", normalization_factory::rms_norm(d_model)),
        (
            "Efficient ScaleNorm",
            normalization_factory::scale_norm(d_model),
        ),
        ("L2 Normalization", normalization_factory::l2_norm(d_model)),
        (
            "Custom Power Norm",
            normalization_factory::power_norm(d_model, 1.5),
        ),
    ];

    for (name, config) in configs {
        info!("  Testing {}...", name);
        let config = config.with_device(0);

        match UnifiedNorm::from_config(&config) {
            Ok(norm) => {
                let output = norm.forward(input)?;
                info!(
                    "    {} - Type: {}, Shape: {:?}, Complexity: {:.2}x",
                    name,
                    norm.norm_type(),
                    output.shape().dims(),
                    config.norm_type.relative_complexity()
                );
                info!("    Use case: {}", config.norm_type.recommended_use_cases());
            }
            Err(e) => warn!("    Failed to create {}: {}", name, e),
        }
    }

    Ok(())
}

/// Demonstrate PreNorm vs PostNorm configurations
async fn demo_norm_positions(input: &Tensor, device: &Device) -> Result<()> {
    info!("\n📍 Demo 3: PreNorm vs PostNorm Configurations");

    let d_model = input.shape().dims()[2];

    // Create a simple mock sublayer for demonstration
    #[derive(Debug)]
    struct MockSublayer {
        device: Device,
    }

    impl Layer for MockSublayer {
        fn forward(&self, input: &Tensor) -> Result<Tensor> {
            // Simple identity + small perturbation
            let perturbation = Tensor::randn(input.shape().clone(), &self.device)? * 0.1;
            Ok((input + &perturbation)?)
        }

        fn layer_type(&self) -> &'static str {
            "MockSublayer"
        }

        fn device(&self) -> &Device {
            &self.device
        }

        fn parameters(&self) -> Vec<&Tensor> {
            vec![]
        }

        fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
            vec![]
        }
    }

    let sublayer = MockSublayer {
        device: device.clone(),
    };

    // Test PreNorm configuration
    info!("  Testing PreNorm (Norm -> Sublayer -> Residual)...");
    let norm_pre = RMSNorm::new(d_model, 1e-6, device)?;
    let prenorm_block = PreNorm::new(norm_pre, sublayer);
    let prenorm_output = prenorm_block.forward(input)?;
    info!(
        "    PreNorm output shape: {:?}",
        prenorm_output.shape().dims()
    );

    // Test PostNorm configuration
    info!("  Testing PostNorm (Sublayer -> Residual -> Norm)...");
    let norm_post = RMSNorm::new(d_model, 1e-6, device)?;
    let sublayer2 = MockSublayer {
        device: device.clone(),
    };
    let postnorm_block = PostNorm::new(norm_post, sublayer2);
    let postnorm_output = postnorm_block.forward(input)?;
    info!(
        "    PostNorm output shape: {:?}",
        postnorm_output.shape().dims()
    );

    // Test generic wrapper
    info!("  Testing NormWrapper with different positions...");
    for position in [NormPosition::Pre, NormPosition::Post] {
        let norm = RMSNorm::new(d_model, 1e-6, device)?;
        let sublayer = MockSublayer {
            device: device.clone(),
        };
        let wrapper = NormWrapper::new(position, norm, sublayer);
        let output = wrapper.forward(input)?;
        info!(
            "    {:?} wrapper output shape: {:?}",
            position,
            output.shape().dims()
        );
    }

    Ok(())
}

/// Demonstrate DeepNorm for very deep networks
async fn demo_deep_norm(input: &Tensor, device: &Device) -> Result<()> {
    info!("\n🏗️  Demo 4: DeepNorm for Very Deep Networks");

    let d_model = input.shape().dims()[2];

    // Test different network depths
    let depths = vec![10, 100, 1000];

    for num_layers in depths {
        info!("  Testing DeepNorm for {} layers...", num_layers);

        let config = DeepNormConfig::new(num_layers, d_model);
        let deep_norm = DeepNorm::new(config, device)?;

        info!("    Alpha (residual scaling): {:.4}", deep_norm.alpha());
        info!("    Beta (init scaling): {:.4}", deep_norm.beta());

        let output = deep_norm.forward(input)?;
        info!(
            "    DeepNorm-{} output shape: {:?}",
            num_layers,
            output.shape().dims()
        );

        // Demonstrate residual connection scaling
        let sublayer_output = input * 0.1; // Mock sublayer output
        let residual_output = deep_norm.forward_residual(input, &sublayer_output)?;
        info!(
            "    Residual output shape: {:?}",
            residual_output.shape().dims()
        );
    }

    Ok(())
}

/// Demonstrate adaptive normalization with conditioning
async fn demo_adaptive_norm(input: &Tensor, device: &Device) -> Result<()> {
    info!("\n🎛️  Demo 5: Adaptive Normalization with Conditioning");

    let d_model = input.shape().dims()[2];
    let condition_dim = 128;
    let batch_size = input.shape().dims()[0];
    let seq_len = input.shape().dims()[1];

    // Create conditioning input
    let condition = Tensor::randn(
        Shape::new(vec![batch_size, seq_len, condition_dim])?,
        device,
    )?;

    info!("  Testing AdaLayerNorm with conditioning...");
    let ada_norm = AdaLayerNorm::new(d_model, condition_dim, 1e-5, device)?;

    // Test with conditioning
    let conditioned_output = ada_norm.forward_conditioned(input, &condition)?;
    info!(
        "    Conditioned output shape: {:?}",
        conditioned_output.shape().dims()
    );

    // Test without conditioning (fallback)
    let unconditioned_output = ada_norm.forward(input)?;
    info!(
        "    Unconditioned output shape: {:?}",
        unconditioned_output.shape().dims()
    );

    // Test through unified framework
    let config = normalization_factory::ada_layer_norm(d_model, condition_dim).with_device(0);

    match UnifiedNorm::from_config(&config) {
        Ok(unified_norm) => {
            let unified_output = unified_norm.forward_conditioned(input, Some(&condition))?;
            info!(
                "    Unified AdaLayerNorm output shape: {:?}",
                unified_output.shape().dims()
            );
        }
        Err(e) => warn!("    Failed to create unified AdaLayerNorm: {}", e),
    }

    Ok(())
}

/// Demonstrate performance comparison between different norms
async fn demo_performance_comparison(input: &Tensor, device: &Device) -> Result<()> {
    info!("\n⚡ Demo 6: Performance and Complexity Comparison");

    let d_model = input.shape().dims()[2];

    let norm_types = vec![
        NormalizationType::LayerNorm {
            elementwise_affine: true,
        },
        NormalizationType::RMSNorm,
        NormalizationType::ScaleNorm,
        NormalizationType::L1Norm,
        NormalizationType::L2Norm,
        NormalizationType::MaxNorm,
        NormalizationType::PowerNorm { power: 1.5 },
    ];

    info!("  Normalization Performance Summary:");
    info!(
        "  {:20} | {:12} | {:8} | {}",
        "Type", "Complexity", "Cond.", "Use Case"
    );
    info!("  {:-<20} | {:-<12} | {:-<8} | {}", "", "", "", "");

    for norm_type in norm_types {
        let complexity = norm_type.relative_complexity();
        let conditioning = if norm_type.requires_conditioning() {
            "Yes"
        } else {
            "No"
        };
        let use_case = norm_type.recommended_use_cases();

        info!(
            "  {:20} | {:8.2}x    | {:8} | {}",
            format!("{:?}", norm_type)
                .split('{')
                .next()
                .unwrap_or("Unknown"),
            complexity,
            conditioning,
            use_case.split(',').next().unwrap_or(use_case)
        );
    }

    // Test actual forward passes for timing (simplified)
    info!("\n  Testing forward pass execution...");
    let configs = vec![
        ("LayerNorm", normalization_factory::layer_norm(d_model)),
        ("RMSNorm", normalization_factory::rms_norm(d_model)),
        ("ScaleNorm", normalization_factory::scale_norm(d_model)),
    ];

    for (name, config) in configs {
        let config = config.with_device(0);
        if let Ok(norm) = UnifiedNorm::from_config(&config) {
            // Simple timing test (not precise but illustrative)
            let start = std::time::Instant::now();
            for _ in 0..10 {
                let _ = norm.forward(input)?;
            }
            let duration = start.elapsed();
            info!("    {} - 10 iterations: {:?}", name, duration);
        }
    }

    Ok(())
}

/// Helper function to show tensor statistics (for debugging/analysis)
#[allow(dead_code)]
fn show_tensor_stats(tensor: &Tensor, name: &str) -> Result<()> {
    let data = tensor.to_vec()?;
    let mean = data.iter().sum::<f32>() / data.len() as f32;
    let variance = data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;
    let std_dev = variance.sqrt();
    let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
    let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));

    info!(
        "    {} stats: mean={:.6}, std={:.6}, min={:.6}, max={:.6}",
        name, mean, std_dev, min_val, max_val
    );
    Ok(())
}