100 lines
3.3 KiB
Rust
100 lines
3.3 KiB
Rust
//! Test to verify correct Device usage in rtx-transformers
|
|
//! RED phase: This test should fail with compilation errors first
|
|
//!
|
|
//! NOTE: Disabled until transformer API is fully implemented
|
|
|
|
#![cfg(feature = "disabled_tests")]
|
|
|
|
use rtx_tensor::{Device, Tensor};
|
|
use rtx_transformers::architectures::TransformerModel;
|
|
use rtx_transformers::architectures::bert::BertModel;
|
|
|
|
#[test]
|
|
fn test_bert_model_device_creation() {
|
|
// RED: This test verifies that BertModel can be created with proper device
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// This should compile and run without issues
|
|
let model = BertModel::new(768, 12, 12, 512, 30522, 512, &device);
|
|
assert!(model.is_ok(), "BertModel should be created successfully");
|
|
|
|
let model = model.unwrap();
|
|
// Verify the model is on the correct device
|
|
assert_eq!(model.device(), &device);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_model_forward_pass() {
|
|
// RED: Test forward pass with proper device handling
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
let model = BertModel::new(768, 12, 12, 512, 30522, 512, &device).unwrap();
|
|
|
|
// Create input tensors on the same device
|
|
let input_ids = Tensor::zeros(&[2, 128], &device).unwrap();
|
|
let attention_mask = Tensor::ones(&[2, 128], &device).unwrap();
|
|
|
|
// Forward pass should work
|
|
let output = model.forward(&input_ids, Some(&attention_mask));
|
|
assert!(output.is_ok(), "Forward pass should succeed");
|
|
|
|
let output = output.unwrap();
|
|
assert_eq!(output.device(), &device, "Output should be on same device");
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_model_parameters_on_device() {
|
|
// RED: Test that all parameters are on the correct device
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
let model = BertModel::new(768, 12, 12, 512, 30522, 512, &device).unwrap();
|
|
|
|
let params = model.parameters();
|
|
assert!(!params.is_empty(), "Model should have parameters");
|
|
|
|
// All parameters should be on the same device
|
|
for param in params {
|
|
assert_eq!(
|
|
param.device(),
|
|
&device,
|
|
"All parameters should be on same device"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_model_gradient_flow() {
|
|
// RED: Test gradient flow with device consistency
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
let mut model = BertModel::new(768, 12, 12, 512, 30522, 512, &device).unwrap();
|
|
model.train(true);
|
|
|
|
// Create input with requires_grad
|
|
let mut input_ids = Tensor::zeros(&[1, 64], &device).unwrap();
|
|
input_ids.set_requires_grad(true).unwrap();
|
|
|
|
let output = model.forward(&input_ids, None).unwrap();
|
|
|
|
// Output should track gradients
|
|
assert!(
|
|
output.requires_grad(),
|
|
"Output should require gradients in training mode"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_transfer() {
|
|
// RED: Test transferring model between devices
|
|
let initial_device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
let mut model = BertModel::new(768, 12, 12, 512, 30522, 512, &initial_device).unwrap();
|
|
|
|
// Try to transfer to same device (should work)
|
|
let result = model.to_device(&initial_device);
|
|
assert!(result.is_ok(), "Transfer to same device should work");
|
|
|
|
// Verify model is still on correct device
|
|
assert_eq!(model.device(), &initial_device);
|
|
}
|