rustytorch/crates/core/rtx-autograd/tests/integration.rs

//! Integration tests for tensor autograd functionality
//!
//! These tests verify that the autograd engine works correctly with tensor operations.
//!
//! Note: These tests use the tape-based API which is not yet implemented in the
//! decorator pattern architecture. They are disabled until the tape API is available.

#![cfg(feature = "tape_api")]

use approx::assert_relative_eq;
use rtx_autograd::{TensorAutograd, clear_tape, enable_grad, tensor_with_grad};
use rtx_tensor::{Device, Tensor};

#[test]
fn test_simple_autograd_integration() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    // Clear any existing tape
    clear_tape();

    enable_grad(|| {
        // Create leaf tensors with gradients
        let a = tensor_with_grad(vec![2.0], [1], &device).unwrap();
        let b = tensor_with_grad(vec![3.0], [1], &device).unwrap();

        // Test that tensors have gradient requirements
        assert_eq!(a.requires_grad(), true);
        assert_eq!(b.requires_grad(), true);

        // Test that node IDs are assigned
        assert!(a.autograd_node_id().is_some());
        assert!(b.autograd_node_id().is_some());

        // Perform operations with autograd recording
        let c = a.add_grad(&b).unwrap();

        // Verify the result has gradients enabled and node ID
        assert_eq!(c.requires_grad(), true);
        assert!(c.autograd_node_id().is_some());

        // Check the computation result
        let c_data = c.to_cpu().unwrap();
        assert_relative_eq!(c_data[0], 5.0, epsilon = 1e-6);
    });
}

#[test]
fn test_multiplication_with_autograd() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    clear_tape();

    enable_grad(|| {
        let a = tensor_with_grad(vec![2.0, 3.0], [2, 1], &device).unwrap();
        let b = tensor_with_grad(vec![4.0, 5.0], [2, 1], &device).unwrap();

        // Test multiplication
        let c = a.mul_grad(&b).unwrap();

        assert_eq!(c.requires_grad(), true);
        assert!(c.autograd_node_id().is_some());

        let c_data = c.to_cpu().unwrap();
        assert_relative_eq!(c_data[0], 8.0, epsilon = 1e-6);
        assert_relative_eq!(c_data[1], 15.0, epsilon = 1e-6);
    });
}

#[test]
fn test_matrix_multiplication_with_autograd() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    clear_tape();

    enable_grad(|| {
        // 2x2 @ 2x2 -> 2x2
        let a = tensor_with_grad(vec![1.0, 2.0, 3.0, 4.0], [2, 2], &device).unwrap();
        let b = tensor_with_grad(vec![5.0, 6.0, 7.0, 8.0], [2, 2], &device).unwrap();

        let c = a.matmul_grad(&b).unwrap();

        assert_eq!(c.requires_grad(), true);
        assert!(c.autograd_node_id().is_some());

        let c_data = c.to_cpu().unwrap();
        // Expected: [[1*5+2*7, 1*6+2*8], [3*5+4*7, 3*6+4*8]] = [[19, 22], [43, 50]]
        assert_relative_eq!(c_data[0], 19.0, epsilon = 1e-6);
        assert_relative_eq!(c_data[1], 22.0, epsilon = 1e-6);
        assert_relative_eq!(c_data[2], 43.0, epsilon = 1e-6);
        assert_relative_eq!(c_data[3], 50.0, epsilon = 1e-6);
    });
}

#[test]
fn test_sum_reduction_with_autograd() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    clear_tape();

    enable_grad(|| {
        let a = tensor_with_grad(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3], &device).unwrap();

        // Sum all elements
        let sum_all = a.sum_grad(None).unwrap();

        assert_eq!(sum_all.requires_grad(), true);
        assert!(sum_all.autograd_node_id().is_some());

        let sum_data = sum_all.to_cpu().unwrap();
        assert_relative_eq!(sum_data[0], 21.0, epsilon = 1e-6);

        // Sum along dimension 0
        let sum_dim0 = a.sum_grad(Some(0)).unwrap();

        assert_eq!(sum_dim0.requires_grad(), true);
        assert!(sum_dim0.autograd_node_id().is_some());

        let sum0_data = sum_dim0.to_cpu().unwrap();
        assert_relative_eq!(sum0_data[0], 5.0, epsilon = 1e-6); // 1+4
        assert_relative_eq!(sum0_data[1], 7.0, epsilon = 1e-6); // 2+5
        assert_relative_eq!(sum0_data[2], 9.0, epsilon = 1e-6); // 3+6
    });
}

#[test]
fn test_complex_computation_graph() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    clear_tape();

    enable_grad(|| {
        // Create a simple neural network-like computation: y = (x * w + b)^2
        let x = tensor_with_grad(vec![1.0, 2.0], [2, 1], &device).unwrap();
        let w = tensor_with_grad(vec![0.5, 1.5], [2, 1], &device).unwrap();
        let b = tensor_with_grad(vec![0.1, 0.2], [2, 1], &device).unwrap();

        // Forward pass
        let linear = x.mul_grad(&w).unwrap();
        let activated = linear.add_grad(&b).unwrap();
        let squared = activated.mul_grad(&activated).unwrap();

        // Verify all intermediate results have gradients
        assert_eq!(linear.requires_grad(), true);
        assert_eq!(activated.requires_grad(), true);
        assert_eq!(squared.requires_grad(), true);

        // Check final result
        let result_data = squared.to_cpu().unwrap();
        assert_relative_eq!(result_data[0], (1.0f32 * 0.5 + 0.1).powi(2), epsilon = 1e-6);
        assert_relative_eq!(result_data[1], (2.0f32 * 1.5 + 0.2).powi(2), epsilon = 1e-6);
    });
}

#[test]
fn test_gradient_storage() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    let a = tensor_with_grad(vec![2.0], [1], &device).unwrap();

    // Initially no gradient
    assert!(a.grad().is_none());

    // Set a gradient
    let grad_tensor = Tensor::ones([1], &device).unwrap();
    a.set_grad(Some(grad_tensor.clone()));

    // Verify gradient is stored
    let stored_grad = a.grad().unwrap();
    let grad_data = stored_grad.to_cpu().unwrap();
    assert_relative_eq!(grad_data[0], 1.0, epsilon = 1e-6);

    // Clear gradient
    a.set_grad(None);
    assert!(a.grad().is_none());
}

#[test]
fn test_mixed_grad_and_non_grad_tensors() {
    let device = Device::cuda(0).unwrap_or(Device::default());

    clear_tape();

    enable_grad(|| {
        // One tensor requires gradients, the other doesn't
        let a = tensor_with_grad(vec![2.0], [1], &device).unwrap();
        let b = Tensor::from_data(vec![3.0], [1], &device).unwrap(); // No gradients

        assert_eq!(a.requires_grad(), true);
        assert_eq!(b.requires_grad(), false);

        // Operation result should require gradients
        let c = a.add_grad(&b).unwrap();

        assert_eq!(c.requires_grad(), true);
        assert!(c.autograd_node_id().is_some());

        let c_data = c.to_cpu().unwrap();
        assert_relative_eq!(c_data[0], 5.0, epsilon = 1e-6);
    });
}