foxhunt/crates/ml/tests/test_grn_weight_initialization.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Test suite for GRN weight initialization verification
//!
//! This test verifies that Gated Residual Network (GRN) layers use proper
//! Xavier/Kaiming weight initialization instead of zeros.
//!
//! Context: Wave 8.6 - Verify that GpuLinear properly initializes
//! weights following Xavier Uniform distribution by default.
//!
//! CRITICAL: Use VarBuilder::from_varmap() for proper weight initialization,
//! NOT VarBuilder::zeros() which creates all-zero weights.

// candle eliminated — test uses native APIs
// candle eliminated — VarMap replaced with GpuVarStore
use std::sync::Arc;
use tracing::info;

use ml::tft::gated_residual::{GRNStack, GatedLinearUnit, GatedResidualNetwork};
use ml::MLError;

/// Calculate mean of a tensor
fn calculate_mean(tensor: &Tensor) -> Result<f32, MLError> {
    let flat = tensor.flatten_all()?;
    let vec = flat.to_vec1::<f32>()?;
    Ok(vec.iter().sum::<f32>() / vec.len() as f32)
}

/// Calculate standard deviation of a tensor
fn calculate_std_dev(tensor: &Tensor) -> Result<f32, MLError> {
    let flat = tensor.flatten_all()?;
    let vec = flat.to_vec1::<f32>()?;
    let mean = vec.iter().sum::<f32>() / vec.len() as f32;
    let variance = vec.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / vec.len() as f32;
    Ok(variance.sqrt())
}

/// Calculate min and max values
fn calculate_range(tensor: &Tensor) -> Result<(f32, f32), MLError> {
    let flat = tensor.flatten_all()?;
    let vec = flat.to_vec1::<f32>()?;
    let min = vec.iter().copied().fold(f32::INFINITY, f32::min);
    let max = vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    Ok((min, max))
}

#[test]
fn test_grn_weight_initialization_statistics() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(64, 64, vs.pp("test"))?;

    // Create test input to extract weight information
    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    // Forward pass to ensure weights are initialized
    let output = grn.forward(&inputs, None)?;

    // Check output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;
    let (min, max) = calculate_range(&output)?;

    info!(mean, std_dev, min, max, "GRN output statistics");

    // Verify non-zero outputs (would be zero if weights were not initialized)
    assert!(
        std_dev > 0.01,
        "Output std dev should be non-zero (got {}), indicating proper weight initialization",
        std_dev
    );

    // Verify output has reasonable range (not all zeros or infinities)
    assert!(
        min.is_finite() && max.is_finite(),
        "Output should be finite"
    );
    assert!(
        (max - min) > 0.1,
        "Output should have non-trivial range (got {})",
        max - min
    );

    Ok(())
}

#[test]
fn test_grn_different_dims_weight_initialization() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Test with different input/output dimensions (triggers skip_projection)
    let grn = GatedResidualNetwork::new(128, 64, vs.pp("test"))?;

    let input_data = vec![1.0f32; 54]; // 2 * 128
    let inputs = Tensor::from_slice(&input_data, (2, 128), &device)?;

    let output = grn.forward(&inputs, None)?;

    // Check output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    info!(mean, std_dev, "GRN (different dims) output statistics");

    // Verify skip projection is also properly initialized
    assert!(
        std_dev > 0.01,
        "Output std dev should be non-zero with skip projection (got {})",
        std_dev
    );

    Ok(())
}

#[test]
fn test_grn_context_projection_initialization() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(64, 64, vs.pp("test"))?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    let context_data = vec![0.5f32; 128]; // 2 * 64
    let context = Tensor::from_slice(&context_data, (2, 64), &device)?;

    // Forward pass with context
    let output_with_context = grn.forward(&inputs, Some(&context))?;

    // Forward pass without context
    let output_no_context = grn.forward(&inputs, None)?;

    // Check that context has an effect (would be same if context_projection not initialized)
    let diff = (output_with_context - output_no_context)?;
    let diff_std = calculate_std_dev(&diff)?;

    info!(diff_std, "Context effect std dev");

    assert!(
        diff_std > 0.01,
        "Context should have measurable effect (got std dev {}), indicating context_projection is initialized",
        diff_std
    );

    Ok(())
}

#[test]
fn test_glu_weight_initialization() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let glu = GatedLinearUnit::new(64, 32, vs.pp("test"))?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    let output = glu.forward(&inputs)?;

    // Check output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    info!(mean, std_dev, "GLU output statistics");

    // GLU uses sigmoid gating, so outputs should be in reasonable range
    assert!(
        std_dev > 0.01,
        "GLU output should have non-zero variance (got {})",
        std_dev
    );

    // Check that output is bounded (sigmoid gate keeps values reasonable)
    let (min, max) = calculate_range(&output)?;
    info!(min, max, "GLU output range");
    assert!(
        min.is_finite() && max.is_finite(),
        "GLU output should be finite"
    );

    Ok(())
}

#[test]
fn test_grn_stack_weight_initialization() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let stack = GRNStack::new(64, 32, 16, 3, vs.pp("test"))?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    let output = stack.forward(&inputs, None)?;

    // Check final output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    info!(mean, std_dev, "GRN stack output statistics");

    // Multi-layer stack should still have non-zero, finite outputs
    assert!(
        std_dev > 0.01,
        "GRN stack output should have non-zero variance (got {})",
        std_dev
    );

    let (min, max) = calculate_range(&output)?;
    assert!(
        min.is_finite() && max.is_finite(),
        "GRN stack output should be finite"
    );

    Ok(())
}

#[test]
fn test_grn_multiple_forward_passes() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Multiple forward passes with different inputs should produce different outputs
    let input1_data = vec![1.0f32; 64]; // 2 * 32
    let input1 = Tensor::from_slice(&input1_data, (2, 32), &device)?;

    let input2_data = vec![2.0f32; 64]; // 2 * 32
    let input2 = Tensor::from_slice(&input2_data, (2, 32), &device)?;

    let output1 = grn.forward(&input1, None)?;
    let output2 = grn.forward(&input2, None)?;

    // Outputs should be different for different inputs
    let diff = (output2 - output1)?;
    let diff_std = calculate_std_dev(&diff)?;

    info!(diff_std, "Output difference std dev");

    assert!(
        diff_std > 0.1,
        "Different inputs should produce different outputs (got std dev {})",
        diff_std
    );

    Ok(())
}

#[test]
fn test_grn_3d_tensor_weight_initialization() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(16, 16, vs.pp("test"))?;

    // 3D input: [batch_size=2, seq_len=5, hidden_dim=16]
    let input_data = vec![1.0f32; 160]; // 2 * 5 * 16
    let inputs = Tensor::from_slice(&input_data, (2, 5, 16), &device)?;

    let output = grn.forward(&inputs, None)?;

    // Check statistics across all dimensions
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    info!(mean, std_dev, "GRN 3D output statistics");

    assert!(
        std_dev > 0.01,
        "3D tensor output should have non-zero variance (got {})",
        std_dev
    );

    Ok(())
}

#[test]
fn test_grn_batch_consistency() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Create two identical samples in a batch
    let mut input_data = vec![1.0f32; 64]; // 2 * 32
                                           // Make second sample different
    for i in 32..64 {
        input_data[i] = 2.0;
    }
    let inputs = Tensor::from_slice(&input_data, (2, 32), &device)?;

    let output = grn.forward(&inputs, None)?;

    // Extract individual batch elements
    let output_vec = output.to_vec2::<f32>()?;
    let sample1 = &output_vec[0];
    let sample2 = &output_vec[1];

    // Calculate difference between samples
    let diff: Vec<f32> = sample1
        .iter()
        .zip(sample2.iter())
        .map(|(a, b)| (a - b).abs())
        .collect();
    let diff_mean = diff.iter().sum::<f32>() / diff.len() as f32;

    info!(diff_mean, "Batch sample difference mean");

    // Different inputs should produce different outputs
    assert!(
        diff_mean > 0.01,
        "Different batch samples should produce different outputs (got mean diff {})",
        diff_mean
    );

    Ok(())
}

#[test]
fn test_grn_zero_input_response() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Zero input
    let zero_input = Tensor::zeros((2, 32), DType::F32, &device)?;
    let output = grn.forward(&zero_input, None)?;

    // Output should not be all zeros if weights are initialized
    // (bias terms and residual connection should produce non-zero output)
    let std_dev = calculate_std_dev(&output)?;

    info!(std_dev, "Zero input output std dev");

    // Note: Even with zero input, properly initialized network should have
    // some non-zero response due to bias terms and layer normalization
    let (min, max) = calculate_range(&output)?;
    assert!(
        min.is_finite() && max.is_finite(),
        "Output should be finite even with zero input"
    );

    Ok(())
}