107 lines
3.4 KiB
Rust
107 lines
3.4 KiB
Rust
/// Tests for tensor operations used in rtx-multimodal
|
|
/// Following strict TDD - RED phase: write failing tests first
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
#[test]
|
|
fn test_tensor_cat_operation() {
|
|
// RED: Test that Tensor::cat works with the correct types
|
|
let device = Device::default();
|
|
|
|
let a = Tensor::ones(&[2, 3, 4], &device).expect("Failed to create tensor a");
|
|
let b = Tensor::ones(&[2, 5, 4], &device).expect("Failed to create tensor b");
|
|
|
|
// Test concatenation along dimension 1
|
|
let result = Tensor::cat(&[a, b], 1);
|
|
assert!(result.is_ok(), "Tensor::cat should work with owned tensors");
|
|
|
|
let concatenated = result.unwrap();
|
|
assert_eq!(concatenated.shape(), &[2, 8, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_cat_with_references() {
|
|
// RED: Test that we can work with tensor references
|
|
let device = Device::default();
|
|
|
|
let a = Tensor::ones(&[2, 3, 4], &device).expect("Failed to create tensor a");
|
|
let b = Tensor::ones(&[2, 5, 4], &device).expect("Failed to create tensor b");
|
|
|
|
// We need to test if cat works with references or if we need to clone
|
|
let a_ref = &a;
|
|
let b_ref = &b;
|
|
|
|
// This might fail if cat doesn't accept references
|
|
let result = Tensor::cat(&[a_ref.clone(), b_ref.clone()], 1);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Should be able to concatenate tensor clones"
|
|
);
|
|
|
|
let concatenated = result.unwrap();
|
|
assert_eq!(concatenated.shape(), &[2, 8, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_async_to_sync_conversion() {
|
|
// RED: Test that we can handle async operations in sync context
|
|
// This test documents the need for a sync wrapper or blocking runtime
|
|
|
|
let device = Device::default();
|
|
let tensor = Tensor::ones(&[2, 3], &device).expect("Failed to create tensor");
|
|
|
|
// We need to test if we can handle async operations
|
|
// For now, this just documents the requirement
|
|
assert!(tensor.shape() == &[2, 3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_narrow_operation() {
|
|
// RED: Test the narrow operation used in gated fusion
|
|
let device = Device::default();
|
|
|
|
let tensor = Tensor::randn(&[2, 6, 4], &device).expect("Failed to create tensor");
|
|
|
|
// Test narrow operation - extract subset along dimension
|
|
let first_half = tensor.narrow(1, 0, 3);
|
|
assert!(first_half.is_ok(), "Narrow should work for first half");
|
|
|
|
let first_half = first_half.unwrap();
|
|
assert_eq!(first_half.shape(), &[2, 3, 4]);
|
|
|
|
let second_half = tensor.narrow(1, 3, 3);
|
|
assert!(second_half.is_ok(), "Narrow should work for second half");
|
|
|
|
let second_half = second_half.unwrap();
|
|
assert_eq!(second_half.shape(), &[2, 3, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_unsqueeze_operation() {
|
|
// RED: Test unsqueeze operation for broadcasting
|
|
let device = Device::default();
|
|
|
|
let tensor = Tensor::ones(&[2, 3], &device).expect("Failed to create tensor");
|
|
|
|
// Test unsqueeze to add dimension
|
|
let unsqueezed = tensor.unsqueeze(-1);
|
|
assert!(unsqueezed.is_ok(), "Unsqueeze should work");
|
|
|
|
let unsqueezed = unsqueezed.unwrap();
|
|
assert_eq!(unsqueezed.shape(), &[2, 3, 1]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_full_operation() {
|
|
// RED: Test creating a tensor filled with a specific value
|
|
let device = Device::default();
|
|
|
|
let shape = vec![2, 3, 4];
|
|
let value = 0.5;
|
|
|
|
let tensor = Tensor::full(&shape, value, &device);
|
|
assert!(tensor.is_ok(), "Should create tensor filled with value");
|
|
|
|
let tensor = tensor.unwrap();
|
|
assert_eq!(tensor.shape(), &[2, 3, 4]);
|
|
}
|