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.

use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaStream};
use ml::cuda_autograd::stream_ops::StreamTensor as GpuTensor;
use ml::tft::gated_residual::{GRNStack, GatedLinearUnit, GatedResidualNetwork};
use ml::MLError;
use tracing::info;

fn test_stream() -> Arc<CudaStream> {
    let ctx = CudaContext::new(0).expect("CUDA required");
    ctx.new_stream().expect("Failed to create CUDA stream")
}

/// Calculate mean of a tensor (downloads to host)
fn calculate_mean(vec: &[f32]) -> f32 {
    vec.iter().sum::<f32>() / vec.len() as f32
}

/// Calculate standard deviation of a tensor (downloads to host)
fn calculate_std_dev(vec: &[f32]) -> f32 {
    let mean = calculate_mean(vec);
    let variance = vec.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / vec.len() as f32;
    variance.sqrt()
}

/// Calculate min and max values
fn calculate_range(vec: &[f32]) -> (f32, f32) {
    let min = vec.iter().copied().fold(f32::INFINITY, f32::min);
    let max = vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    (min, max)
}

#[test]
fn test_grn_weight_initialization_statistics() -> Result<(), MLError> {
    let stream = test_stream();

    let grn = GatedResidualNetwork::new(64, 64, &stream)?;

    // Create test input
    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = GpuTensor::from_vec(input_data, &[2, 64], &stream)?;

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

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

    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 stream = test_stream();

    // Test with different input/output dimensions (triggers skip_projection)
    let grn = GatedResidualNetwork::new(128, 64, &stream)?;

    let input_data = vec![1.0f32; 256]; // 2 * 128
    let inputs = GpuTensor::from_vec(input_data, &[2, 128], &stream)?;

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

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

    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 stream = test_stream();

    let grn = GatedResidualNetwork::new(64, 64, &stream)?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = GpuTensor::from_vec(input_data, &[2, 64], &stream)?;

    let context_data = vec![0.5f32; 128]; // 2 * 64
    let context = GpuTensor::from_vec(context_data, &[2, 64], &stream)?;

    // 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 with_ctx_vec = output_with_context.to_vec()?;
    let no_ctx_vec = output_no_context.to_vec()?;
    let diff_vec: Vec<f32> = with_ctx_vec
        .iter()
        .zip(no_ctx_vec.iter())
        .map(|(a, b)| a - b)
        .collect();
    let diff_std = calculate_std_dev(&diff_vec);

    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 stream = test_stream();

    let glu = GatedLinearUnit::new(64, 32, &stream)?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = GpuTensor::from_vec(input_data, &[2, 64], &stream)?;

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

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

    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_vec);
    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 stream = test_stream();

    let stack = GRNStack::new(64, 32, 16, 3, &stream)?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = GpuTensor::from_vec(input_data, &[2, 64], &stream)?;

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

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

    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_vec);
    assert!(
        min.is_finite() && max.is_finite(),
        "GRN stack output should be finite"
    );

    Ok(())
}

#[test]
fn test_grn_multiple_forward_passes() -> Result<(), MLError> {
    let stream = test_stream();

    let grn = GatedResidualNetwork::new(32, 32, &stream)?;

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

    let input2_data = vec![2.0f32; 64]; // 2 * 32
    let input2 = GpuTensor::from_vec(input2_data, &[2, 32], &stream)?;

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

    // Outputs should be different for different inputs
    let out1_vec = output1.to_vec()?;
    let out2_vec = output2.to_vec()?;
    let diff_vec: Vec<f32> = out2_vec
        .iter()
        .zip(out1_vec.iter())
        .map(|(a, b)| a - b)
        .collect();
    let diff_std = calculate_std_dev(&diff_vec);

    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 stream = test_stream();

    let grn = GatedResidualNetwork::new(16, 16, &stream)?;

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

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

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

    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 stream = test_stream();

    let grn = GatedResidualNetwork::new(32, 32, &stream)?;

    // 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 = GpuTensor::from_vec(input_data, &[2, 32], &stream)?;

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

    // Extract individual batch elements (flat vector, rows of 32)
    let output_vec = output.to_vec()?;
    let sample1 = &output_vec[..32];
    let sample2 = &output_vec[32..];

    // 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 stream = test_stream();

    let grn = GatedResidualNetwork::new(32, 32, &stream)?;

    // Zero input
    let zero_input = GpuTensor::zeros(&[2, 32], &stream)?;
    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 output_vec = output.to_vec()?;
    let std_dev = calculate_std_dev(&output_vec);

    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_vec);
    assert!(
        min.is_finite() && max.is_finite(),
        "Output should be finite even with zero input"
    );

    Ok(())
}