foxhunt/crates/ml/tests/mamba2_gradient_extraction_test.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,
)]
//! # MAMBA-2 Gradient Extraction Test (TDD)
//!
//! **Test-Driven Development**: This test verifies that gradients are properly extracted
//! from VarMap parameters after backward() pass.
//!
//! **Root Cause**: backward_pass() uses zeros_like() placeholder gradients instead of
//! extracting real gradients from VarMap.
//!
//! **Expected Behavior**:
//! 1. Call forward() to compute loss
//! 2. Call backward() to compute gradients
//! 3. Extract gradients from VarMap parameters (input_proj, output_proj, layer_norms)
//! 4. Verify gradients are non-zero and valid (not NaN/Inf)

#![allow(unused_crate_dependencies)]

use candle_core::{DType, Device, Tensor};
use ml::mamba::Mamba2SSM;
use ml::MLError;
use tracing::info;

#[test]
fn test_mamba2_gradient_extraction_from_varmap() -> Result<(), MLError> {
    info!("=== MAMBA-2 Gradient Extraction Test ===");

    let device = Device::cuda_if_available(0)?;
    info!(?device, "Device");

    // Create small MAMBA-2 model
    let mut model = Mamba2SSM::default_hft(&device)?;
    info!(num_parameters = model.metadata.num_parameters, "Model created");

    // Create dummy input and target
    let batch_size = model.config.batch_size;
    let seq_len = model.config.seq_len;
    let d_model = model.config.d_model;

    let input_data = vec![0.1f64; batch_size * seq_len * d_model];
    let input = Tensor::from_vec(input_data, (batch_size, seq_len, d_model), &device)?;

    let target_data = vec![0.5f64; batch_size * seq_len];
    let target = Tensor::from_vec(target_data, (batch_size, seq_len, 1), &device)?;

    info!(input_dims = ?input.dims(), "Input shape");
    info!(target_dims = ?target.dims(), "Target shape");

    // Forward pass
    let output = model.forward(&input)?;
    info!(output_dims = ?output.dims(), "Output shape");

    // Compute loss (MSE)
    let diff = output.broadcast_sub(&target)?;
    let loss = diff.sqr()?.mean_all()?;
    let loss_value = loss.to_scalar::<f64>()?;
    info!(loss_value, "Loss");

    // Backward pass - THIS SHOULD COMPUTE REAL GRADIENTS
    let grads = loss.backward()?;

    // Extract gradients from GradStore
    info!("=== Extracting Gradients from GradStore ===");

    let varmap = &model.varmap;
    let all_vars = varmap.all_vars();
    info!(total_vars = all_vars.len(), "Total VarMap variables");

    let mut vars_with_gradients = 0;
    let mut total_grad_norm = 0.0f64;

    for (idx, var) in all_vars.iter().enumerate() {
        if let Some(grad) = grads.get(var) {
            // Compute gradient norm
            let grad_vec = grad.flatten_all()?.to_vec1::<f64>()?;
            let grad_norm: f64 = grad_vec.iter().map(|&g| g.powi(2)).sum::<f64>().sqrt();

            info!(var_idx = idx, grad_norm, "Var gradient norm");

            // Verify gradient is valid
            assert!(!grad_norm.is_nan(), "Gradient {} is NaN", idx);
            assert!(!grad_norm.is_infinite(), "Gradient {} is Inf", idx);

            if grad_norm > 1e-9 {
                vars_with_gradients += 1;
                total_grad_norm += grad_norm;
            }
        } else {
            info!(var_idx = idx, "Var has no gradient");
        }
    }

    info!(
        vars_with_gradients,
        total_vars = all_vars.len(),
        total_grad_norm,
        "Gradient summary"
    );

    // CRITICAL ASSERTION: At least some parameters should have non-zero gradients
    assert!(
        vars_with_gradients > 0,
        "FAIL: No variables have gradients! backward() did not compute gradients."
    );

    assert!(
        total_grad_norm > 1e-6,
        "FAIL: Total gradient norm is too small ({:.6}). Gradients may be zeros.",
        total_grad_norm
    );

    info!("TEST PASSED: Gradients extracted from VarMap");
    Ok(())
}

#[test]
fn test_mamba2_backward_pass_extracts_real_gradients() -> Result<(), MLError> {
    info!("=== MAMBA-2 backward_pass() Real Gradient Test ===");

    let device = Device::cuda_if_available(0)?;
    let mut model = Mamba2SSM::default_hft(&device)?;

    // Create input/target
    let batch_size = model.config.batch_size;
    let seq_len = model.config.seq_len;
    let d_model = model.config.d_model;

    let input = Tensor::ones((batch_size, seq_len, d_model), DType::F64, &device)?;
    let target = Tensor::ones((batch_size, seq_len, 1), DType::F64, &device)?;

    // Forward + loss
    let output = model.forward(&input)?;
    let diff = output.broadcast_sub(&target)?;
    let loss = diff.sqr()?.mean_all()?;

    info!(loss = loss.to_scalar::<f64>()?, "Loss");

    // Call backward_pass (current implementation uses zeros_like placeholders)
    model.backward_pass(&loss, &input, &target)?;

    // Check model.gradients HashMap
    info!(total_entries = model.gradients.len(), "=== Model Gradients HashMap ===");

    for (key, grad) in model.gradients.iter() {
        let grad_vec = grad.flatten_all()?.to_vec1::<f64>()?;
        let grad_norm: f64 = grad_vec.iter().map(|&g| g.powi(2)).sum::<f64>().sqrt();
        let is_nonzero = grad_norm > 1e-9;
        info!(key, grad_norm, is_nonzero, "Gradient entry");
    }

    // This test will FAIL until we fix backward_pass()
    // After fix, gradients should be non-zero
    let total_grad_norm: f64 = model
        .gradients
        .values()
        .map(|grad| {
            let grad_vec = grad.flatten_all().unwrap().to_vec1::<f64>().unwrap();
            grad_vec.iter().map(|&g| g.powi(2)).sum::<f64>().sqrt()
        })
        .sum();

    info!(total_grad_norm, "Total gradient norm in model.gradients");

    // EXPECTED TO FAIL with current zeros_like implementation
    assert!(
        total_grad_norm > 1e-6,
        "FAIL: backward_pass() produced zero gradients. Need to extract from VarMap."
    );

    info!("TEST PASSED: backward_pass() extracts real gradients");
    Ok(())
}