//! Core Tensor implementation with GPU memory backing //! //! This module provides the main `Tensor` type that serves as the primary interface //! for tensor operations in RustyTorch++. The design emphasizes: //! //! - GPU-first memory management via rtx-runtime //! - Zero-copy views and broadcasting //! - Type safety with compile-time shape checking where possible //! - PyTorch-compatible API surface for easy migration use crate::{Device, DType, Shape, Storage, TensorError, Result}; use std::sync::{Arc, Mutex}; use rand::Rng; /// Simple node ID for autograd integration (temporary until proper integration) #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct NodeId(pub usize); impl NodeId { pub fn new(id: usize) -> Self { NodeId(id) } } /// Core tensor type with GPU memory backing and PyTorch-compatible API /// /// # Design Principles /// - GPU memory backed by rtx-runtime allocator /// - Reference counting for efficient memory usage /// - Zero-copy views via shared storage with offset/stride /// - Automatic gradient computation support (Phase 2) #[derive(Debug, Clone)] pub struct Tensor { /// Underlying storage with GPU memory storage: Arc, /// Tensor shape (dimensions) shape: Shape, /// Memory stride for each dimension strides: Vec, /// Offset into storage buffer offset: usize, /// Data type dtype: DType, /// Device location device: Device, /// Whether gradient computation is required requires_grad: bool, /// Gradient tensor (lazily allocated) grad: Arc>>, /// Node ID in computation graph for autograd node_id: Option, } impl Tensor { /// Create a new tensor with zeros /// /// # Arguments /// * `shape` - Tensor dimensions /// * `device` - Target device for memory allocation /// /// # Returns /// New zero-initialized tensor pub fn zeros>(shape: S, device: &Device) -> Result { let shape = shape.into(); let numel = shape.numel(); let storage = Arc::new(Storage::zeros(numel, DType::F32, device)?); let strides = shape.default_strides(); Ok(Tensor { storage, shape, strides, offset: 0, dtype: DType::F32, device: device.clone(), requires_grad: false, grad: Arc::new(Mutex::new(None)), node_id: None, }) } /// Create a new tensor with ones pub fn ones>(shape: S, device: &Device) -> Result { let shape = shape.into(); let numel = shape.numel(); let storage = Arc::new(Storage::full(numel, 1.0, DType::F32, device)?); let strides = shape.default_strides(); Ok(Tensor { storage, shape, strides, offset: 0, dtype: DType::F32, device: device.clone(), requires_grad: false, grad: Arc::new(Mutex::new(None)), node_id: None, }) } /// Create an identity matrix pub fn eye(n: usize, device: &Device) -> Result { let mut data = vec![0.0f32; n * n]; for i in 0..n { data[i * n + i] = 1.0; } let storage = Arc::new(Storage::from_data(data, DType::F32, device)?); let shape = Shape::from(vec![n, n]); let strides = shape.default_strides(); Ok(Tensor { storage, shape, strides, offset: 0, dtype: DType::F32, device: device.clone(), requires_grad: false, grad: Arc::new(Mutex::new(None)), node_id: None, }) } /// Create a tensor from raw data pub fn from_data>( data: Vec, shape: S, device: &Device, ) -> Result { let shape = shape.into(); if data.len() != shape.numel() { return Err(TensorError::shape(format!( "Data length {} doesn't match shape size {}", data.len(), shape.numel() ))); } let storage = Arc::new(Storage::from_vec(data, DType::F32, device)?); let strides = shape.default_strides(); Ok(Tensor { storage, shape, strides, offset: 0, dtype: DType::F32, device: device.clone(), requires_grad: false, grad: Arc::new(Mutex::new(None)), node_id: None, }) } /// Get tensor shape pub fn shape(&self) -> &Shape { &self.shape } /// Get number of dimensions pub fn ndim(&self) -> usize { self.shape.ndim() } /// Get total number of elements pub fn numel(&self) -> usize { self.shape.numel() } /// Get device location pub fn device(&self) -> &Device { &self.device } /// Get data type pub fn dtype(&self) -> DType { self.dtype } /// Check if gradients are required pub fn requires_grad(&self) -> bool { self.requires_grad } /// Set gradient requirement pub fn set_requires_grad(&mut self, requires_grad: bool) { self.requires_grad = requires_grad; } /// Get the gradient tensor pub fn grad(&self) -> Option { self.grad.lock().unwrap().clone() } /// Set the gradient tensor pub fn set_grad(&self, grad: Option) { *self.grad.lock().unwrap() = grad; } /// Get the node ID for autograd pub fn node_id(&self) -> Option { self.node_id } /// Set the node ID for autograd (internal use) fn set_node_id(&mut self, node_id: NodeId) { self.node_id = Some(node_id); } /// Create a tensor that requires gradients and has a node ID pub fn require_grad(mut self) -> Self { self.requires_grad = true; // Node ID will be set when added to tape self } /// Get or create a node ID for this tensor (simplified for now) fn get_or_create_node_id(&mut self) -> Result { if let Some(node_id) = self.node_id { return Ok(node_id); } // For now, create a simple incrementing node ID static mut NEXT_NODE_ID: usize = 0; let node_id = unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }; self.node_id = Some(node_id); Ok(node_id) } /// Compute gradients via backward pass starting from this tensor (placeholder) pub fn backward(&self) -> Result<()> { if self.node_id.is_none() { return Err(TensorError::autograd("Cannot backward: tensor has no node ID".to_string())); } // This is a placeholder implementation // The actual backward pass will be implemented in the autograd integration layer // For now, just return success to allow compilation Ok(()) } /// Create a view of this tensor with new shape pub fn view>(&self, shape: S) -> Result { let new_shape = shape.into(); // Check that the new shape has the same number of elements if !self.shape.can_reshape_to(&new_shape) { return Err(TensorError::shape(format!( "Cannot view tensor with {} elements as shape with {} elements", self.shape.numel(), new_shape.numel() ))); } // Check if tensor is contiguous for zero-copy view // For simplicity, we assume contiguous layout for now // In practice, this would check strides vs. default strides let new_strides = new_shape.default_strides(); // Create view sharing the same storage Ok(Tensor { storage: self.storage.clone(), shape: new_shape, strides: new_strides, offset: self.offset, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad, grad: Arc::new(Mutex::new(None)), // New tensor starts with no gradient node_id: None, // Views get new node IDs when used in operations }) } /// Reshape tensor (may require copy if not contiguous) pub fn reshape>(&self, shape: S) -> Result { let new_shape = shape.into(); // Check that the new shape has the same number of elements if !self.shape.can_reshape_to(&new_shape) { return Err(TensorError::shape(format!( "Cannot reshape tensor with {} elements to shape with {} elements", self.shape.numel(), new_shape.numel() ))); } // For now, always create a copy since we're not tracking contiguity // In practice, this would try view() first and fall back to copy if needed let new_numel = new_shape.numel(); let new_storage = Arc::new(Storage::zeros(new_numel, self.dtype, &self.device)?); let new_strides = new_shape.default_strides(); // Copy data from current tensor let data = self.to_cpu()?; let mut result_tensor = Tensor { storage: new_storage, shape: new_shape, strides: new_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad, grad: Arc::new(Mutex::new(None)), // New tensor starts with no gradient node_id: None, // Reshaping creates new node ID when used in operations }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&data)?; Ok(result_tensor) } /// Element-wise addition with autograd support pub fn add(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot add tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Check dtype compatibility if self.dtype != other.dtype { return Err(TensorError::type_error(format!( "Cannot add tensors with different dtypes: {:?} and {:?}", self.dtype, other.dtype ))); } // Compute broadcast shape let result_shape = self.shape.broadcast_with(&other.shape)?; let result_numel = result_shape.numel(); // Create result tensor let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Perform element-wise addition with broadcasting if self.shape == other.shape && self.shape == result_shape { // No broadcasting needed - simple element-wise addition for i in 0..result_numel { result_data[i] = self_data[i] + other_data[i]; } } else { // Broadcasting required for result_idx in 0..result_numel { let result_coords = Self::unravel_index(result_idx, &result_shape); let self_idx = Self::broadcast_index(&result_coords, &self.shape, &result_shape); let other_idx = Self::broadcast_index(&result_coords, &other.shape, &result_shape); result_data[result_idx] = self_data[self_idx] + other_data[other_idx]; } } // Copy result data back to storage let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; // TODO: Record operation in autograd tape - will be implemented in integration layer if self.requires_grad || other.requires_grad { // For now, just assign a node ID to the result tensor static mut NEXT_NODE_ID: usize = 100; // Start from 100 to avoid conflicts result_tensor.node_id = Some(unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }); } Ok(result_tensor) } /// Element-wise multiplication pub fn mul(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot multiply tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Check dtype compatibility if self.dtype != other.dtype { return Err(TensorError::type_error(format!( "Cannot multiply tensors with different dtypes: {:?} and {:?}", self.dtype, other.dtype ))); } // Compute broadcast shape let result_shape = self.shape.broadcast_with(&other.shape)?; let result_numel = result_shape.numel(); // Create result tensor let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Perform element-wise multiplication with broadcasting if self.shape == other.shape && self.shape == result_shape { // No broadcasting needed - simple element-wise multiplication for i in 0..result_numel { result_data[i] = self_data[i] * other_data[i]; } } else { // Broadcasting required for result_idx in 0..result_numel { let result_coords = Self::unravel_index(result_idx, &result_shape); let self_idx = Self::broadcast_index(&result_coords, &self.shape, &result_shape); let other_idx = Self::broadcast_index(&result_coords, &other.shape, &result_shape); result_data[result_idx] = self_data[self_idx] * other_data[other_idx]; } } // Copy result data back to storage let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; // TODO: Record operation in autograd tape - will be implemented in integration layer if self.requires_grad || other.requires_grad { // For now, just assign a node ID to the result tensor static mut NEXT_NODE_ID: usize = 200; // Start from 200 to avoid conflicts result_tensor.node_id = Some(unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }); } Ok(result_tensor) } /// Matrix multiplication pub fn matmul(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot matmul tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Check dtype compatibility if self.dtype != other.dtype { return Err(TensorError::type_error(format!( "Cannot matmul tensors with different dtypes: {:?} and {:?}", self.dtype, other.dtype ))); } // Compute result shape and validate matmul compatibility let result_shape = self.shape.matmul_shape(&other.shape)?; let result_numel = result_shape.numel(); // Create result tensor let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // For 2D matrix multiplication: (m, k) @ (k, n) -> (m, n) if self.shape.ndim() == 2 && other.shape.ndim() == 2 { let m = self.shape.dims()[0]; let k = self.shape.dims()[1]; let n = other.shape.dims()[1]; for i in 0..m { for j in 0..n { let mut sum = 0.0; for kk in 0..k { let self_idx = i * k + kk; let other_idx = kk * n + j; sum += self_data[self_idx] * other_data[other_idx]; } let result_idx = i * n + j; result_data[result_idx] = sum; } } } else { // For higher dimensions, handle batch dimensions // This is a simplified implementation - in practice would need more sophisticated handling return Err(TensorError::shape( "Batch matrix multiplication not yet implemented" )); } // Copy result data back to storage let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; // TODO: Record operation in autograd tape - will be implemented in integration layer if self.requires_grad || other.requires_grad { // For now, just assign a node ID to the result tensor static mut NEXT_NODE_ID: usize = 300; // Start from 300 to avoid conflicts result_tensor.node_id = Some(unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }); } Ok(result_tensor) } /// Sum reduction along specified dimension pub fn sum(&self, dim: Option) -> Result { // Compute result shape let result_shape = self.shape.sum_shape(dim)?; let result_numel = result_shape.numel(); // Create result tensor let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; match dim { None => { // Sum all elements -> scalar result let sum: f32 = self_data.iter().sum(); result_data[0] = sum; } Some(axis) => { // Sum along specific dimension let self_dims = self.shape.dims(); if axis >= self_dims.len() { return Err(TensorError::shape(format!( "Dimension {} out of range for tensor with {} dimensions", axis, self_dims.len() ))); } // Calculate strides and perform reduction if self_dims.len() == 2 { // Handle 2D case efficiently match axis { 0 => { // Sum along rows -> result has shape [cols] let rows = self_dims[0]; let cols = self_dims[1]; for j in 0..cols { let mut sum = 0.0; for i in 0..rows { sum += self_data[i * cols + j]; } result_data[j] = sum; } } 1 => { // Sum along cols -> result has shape [rows] let rows = self_dims[0]; let cols = self_dims[1]; for i in 0..rows { let mut sum = 0.0; for j in 0..cols { sum += self_data[i * cols + j]; } result_data[i] = sum; } } _ => unreachable!("axis checked above"), } } else { // General n-dimensional case // For simplicity, use a more general but less efficient approach let self_strides = self.shape.default_strides(); let result_dims = result_shape.dims(); for result_idx in 0..result_numel { let mut result_coords = Self::unravel_index(result_idx, &result_shape); // Insert the summed dimension result_coords.insert(axis, 0); let mut sum = 0.0; for k in 0..self_dims[axis] { result_coords[axis] = k; let self_idx = Self::ravel_index(&result_coords, &self_strides); sum += self_data[self_idx]; } result_data[result_idx] = sum; } } } } // Copy result data back to storage let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; // TODO: Record operation in autograd tape - will be implemented in integration layer if self.requires_grad { // For now, just assign a node ID to the result tensor static mut NEXT_NODE_ID: usize = 400; // Start from 400 to avoid conflicts result_tensor.node_id = Some(unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }); } Ok(result_tensor) } /// Copy tensor data to CPU as Vec pub fn to_cpu(&self) -> Result> { self.storage.to_cpu() } /// Move tensor to specified device pub fn to_device(&self, device: &Device) -> Result { if device == &self.device { return Ok(self.clone()); } let new_storage = Arc::new(self.storage.to_device(device)?); Ok(Tensor { storage: new_storage, shape: self.shape.clone(), strides: self.strides.clone(), offset: self.offset, dtype: self.dtype, device: device.clone(), requires_grad: self.requires_grad, grad: Arc::new(Mutex::new(None)), // Device transfer creates new tensor with no grad node_id: None, // Device transfer creates new node ID }) } /// Convert flat index to multi-dimensional coordinates fn unravel_index(index: usize, shape: &Shape) -> Vec { let dims = shape.dims(); let mut coords = Vec::with_capacity(dims.len()); let mut remaining = index; for &dim_size in dims { coords.push(remaining % dim_size); remaining /= dim_size; } coords.reverse(); coords } /// Convert multi-dimensional coordinates to flat index for broadcasting fn broadcast_index(coords: &[usize], shape: &Shape, result_shape: &Shape) -> usize { let dims = shape.dims(); let result_dims = result_shape.dims(); let ndim_diff = result_dims.len() - dims.len(); let mut index = 0; let mut stride = 1; // Iterate from right to left (reverse order) for i in (0..dims.len()).rev() { let coord_idx = ndim_diff + i; let coord = if coord_idx < coords.len() { if dims[i] == 1 { 0 // Broadcast dimension } else { coords[coord_idx] } } else { 0 }; index += coord * stride; stride *= dims[i]; } index } /// Convert multi-dimensional coordinates to flat index fn ravel_index(coords: &[usize], strides: &[usize]) -> usize { coords.iter().zip(strides.iter()).map(|(&c, &s)| c * s).sum() } /// Element-wise division pub fn div(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot divide tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Check dtype compatibility if self.dtype != other.dtype { return Err(TensorError::type_error(format!( "Cannot divide tensors with different dtypes: {:?} and {:?}", self.dtype, other.dtype ))); } // Broadcast shapes and create result tensor let result_shape = self.shape.broadcast_with(&other.shape)?; let result_numel = result_shape.numel(); let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Perform element-wise division with broadcasting for i in 0..result_numel { let coords = Self::unravel_index(i, &result_shape); let self_idx = Self::broadcast_index(&coords, &self.shape, &result_shape); let other_idx = Self::broadcast_index(&coords, &other.shape, &result_shape); let other_val = other_data[other_idx]; if other_val.abs() < 1e-8 { return Err(TensorError::invalid_op("Division by zero encountered".to_string())); } result_data[i] = self_data[self_idx] / other_val; } // Create result tensor and copy data let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; // Copy result data back to storage Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; Ok(result_tensor) } /// Element-wise maximum pub fn maximum(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot compute maximum of tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Broadcast shapes and create result tensor let result_shape = self.shape.broadcast_with(&other.shape)?; let result_numel = result_shape.numel(); let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Perform element-wise maximum with broadcasting for i in 0..result_numel { let coords = Self::unravel_index(i, &result_shape); let self_idx = Self::broadcast_index(&coords, &self.shape, &result_shape); let other_idx = Self::broadcast_index(&coords, &other.shape, &result_shape); result_data[i] = self_data[self_idx].max(other_data[other_idx]); } // Create result tensor and copy data let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; // Copy result data back to storage Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; Ok(result_tensor) } /// Element-wise minimum pub fn minimum(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot compute minimum of tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Broadcast shapes and create result tensor let result_shape = self.shape.broadcast_with(&other.shape)?; let result_numel = result_shape.numel(); let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Perform element-wise minimum with broadcasting for i in 0..result_numel { let coords = Self::unravel_index(i, &result_shape); let self_idx = Self::broadcast_index(&coords, &self.shape, &result_shape); let other_idx = Self::broadcast_index(&coords, &other.shape, &result_shape); result_data[i] = self_data[self_idx].min(other_data[other_idx]); } // Create result tensor and copy data let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; // Copy result data back to storage Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; Ok(result_tensor) } /// ReLU activation function pub fn relu(&self) -> Result { let result_numel = self.numel(); let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = self.shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Apply ReLU: max(0, x) for (i, &value) in self_data.iter().enumerate() { result_data[i] = value.max(0.0); } // Create result tensor and copy data let mut result_tensor = Tensor { storage: result_storage, shape: self.shape.clone(), strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; // Copy result data back to storage Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; Ok(result_tensor) } /// Create a new tensor from data and shape (convenience constructor) pub fn new(data: Vec, shape: Vec) -> Result { let device = Device::cpu(); // Default to CPU Self::from_data(data, shape, &device) } /// Get the raw data as a Vec (alias for to_cpu for API compatibility) pub fn data(&self) -> Result> { self.to_cpu() } /// Element-wise subtraction pub fn subtract(&self, other: &Self) -> Result { // Check device compatibility if self.device != other.device { return Err(TensorError::device(format!( "Cannot subtract tensors on different devices: {:?} and {:?}", self.device, other.device ))); } // Check dtype compatibility if self.dtype != other.dtype { return Err(TensorError::type_error(format!( "Cannot subtract tensors with different dtypes: {:?} and {:?}", self.dtype, other.dtype ))); } // Compute broadcast shape let result_shape = self.shape.broadcast_with(&other.shape)?; let result_numel = result_shape.numel(); let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = result_shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let other_data = other.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Perform element-wise subtraction with broadcasting if self.shape == other.shape && self.shape == result_shape { // No broadcasting needed - simple element-wise subtraction for i in 0..result_numel { result_data[i] = self_data[i] - other_data[i]; } } else { // Broadcasting required for result_idx in 0..result_numel { let result_coords = Self::unravel_index(result_idx, &result_shape); let self_idx = Self::broadcast_index(&result_coords, &self.shape, &result_shape); let other_idx = Self::broadcast_index(&result_coords, &other.shape, &result_shape); result_data[result_idx] = self_data[self_idx] - other_data[other_idx]; } } // Copy result data back to storage let mut result_tensor = Tensor { storage: result_storage, shape: result_shape, strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad || other.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; // TODO: Record operation in autograd tape - will be implemented in integration layer if self.requires_grad || other.requires_grad { // For now, just assign a node ID to the result tensor static mut NEXT_NODE_ID: usize = 200; // Start from 200 to avoid conflicts result_tensor.node_id = Some(unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }); } Ok(result_tensor) } /// Scalar multiplication pub fn scalar_mul(&self, scalar: f32) -> Result { let result_numel = self.numel(); let result_storage = Arc::new(Storage::zeros(result_numel, self.dtype, &self.device)?); let result_strides = self.shape.default_strides(); // Get CPU data for computation let self_data = self.to_cpu()?; let mut result_data = vec![0.0f32; result_numel]; // Apply scalar multiplication for (i, &value) in self_data.iter().enumerate() { result_data[i] = value * scalar; } // Create result tensor and copy data let mut result_tensor = Tensor { storage: result_storage, shape: self.shape.clone(), strides: result_strides, offset: 0, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad, grad: Arc::new(Mutex::new(None)), node_id: None, }; Arc::get_mut(&mut result_tensor.storage) .unwrap() .copy_from_cpu(&result_data)?; // TODO: Record operation in autograd tape if self.requires_grad { static mut NEXT_NODE_ID: usize = 300; // Start from 300 to avoid conflicts result_tensor.node_id = Some(unsafe { let id = NEXT_NODE_ID; NEXT_NODE_ID += 1; NodeId::new(id) }); } Ok(result_tensor) } // Additional methods needed for multimodal implementation pub fn transpose(&self, dim0: i32, dim1: i32) -> Result { // Simple transpose implementation for 2D tensors if self.ndim() != 2 { return Ok(self.clone()); // Placeholder - just return self for now } let shape = self.shape(); let dims = shape.dims(); let new_shape = vec![dims[1], dims[0]]; let data = self.to_cpu()?; let mut transposed_data = vec![0.0; data.len()]; for i in 0..dims[0] { for j in 0..dims[1] { transposed_data[j * dims[0] + i] = data[i * dims[1] + j]; } } Tensor::from_data(transposed_data, new_shape, &self.device) } pub fn permute(&self, dims: &[usize]) -> Result { if dims.len() != self.ndim() { return Err(TensorError::shape("Permute dimensions must match tensor dimensions".to_string())); } // Validate permutation dimensions let mut sorted_dims = dims.to_vec(); sorted_dims.sort(); for (i, &dim) in sorted_dims.iter().enumerate() { if dim != i { return Err(TensorError::shape("Invalid permutation dimensions".to_string())); } } // Create new shape and strides based on permutation let old_shape = self.shape().dims(); let old_strides = &self.strides; let mut new_shape = vec![0; dims.len()]; let mut new_strides = vec![0; dims.len()]; for (i, &dim) in dims.iter().enumerate() { new_shape[i] = old_shape[dim]; new_strides[i] = old_strides[dim]; } Ok(Tensor { storage: self.storage.clone(), shape: Shape::from(new_shape), strides: new_strides, offset: self.offset, dtype: self.dtype, device: self.device.clone(), requires_grad: self.requires_grad, grad: self.grad.clone(), node_id: self.node_id, }) } pub fn unsqueeze(&self, dim: i32) -> Result { let mut new_shape = self.shape().dims().to_vec(); let actual_dim = if dim < 0 { (new_shape.len() as i32 + dim + 1) as usize } else { dim as usize }; new_shape.insert(actual_dim, 1); self.view(new_shape) } pub fn squeeze(&self, dim: Option) -> Result { let mut new_shape = self.shape().dims().to_vec(); if let Some(d) = dim { let actual_dim = if d < 0 { (new_shape.len() as i32 + d) as usize } else { d as usize }; if actual_dim < new_shape.len() && new_shape[actual_dim] == 1 { new_shape.remove(actual_dim); } } else { new_shape.retain(|&x| x != 1); } self.view(new_shape) } pub fn expand(&self, shape: &[usize]) -> Result { // Simplified expand - just return reshaped tensor if shape.iter().product::() == self.numel() { self.reshape(shape.to_vec()) } else { Ok(self.clone()) } } pub fn concat(tensors: &[&Self], dim: usize) -> Result { if tensors.is_empty() { return Err(TensorError::shape("Cannot concat empty tensor list".to_string())); } let first_tensor = tensors[0]; let first_shape = first_tensor.shape().dims(); if dim >= first_shape.len() { return Err(TensorError::shape("Concatenation dimension out of bounds".to_string())); } // Validate all tensors have compatible shapes let mut total_dim_size = first_shape[dim]; for (i, tensor) in tensors.iter().enumerate().skip(1) { let shape = tensor.shape().dims(); if shape.len() != first_shape.len() { return Err(TensorError::shape(format!("Tensor {} has incompatible dimensions", i))); } for (j, (&s1, &s2)) in first_shape.iter().zip(shape.iter()).enumerate() { if j != dim && s1 != s2 { return Err(TensorError::shape(format!("Tensor {} has incompatible shape at dimension {}", i, j))); } } total_dim_size += shape[dim]; } // Create result shape let mut result_shape = first_shape.to_vec(); result_shape[dim] = total_dim_size; let result_numel = result_shape.iter().product(); // Create result tensor with all data let mut result_data = Vec::with_capacity(result_numel); // For now, implement a simple CPU-based concatenation // TODO: Use GPU kernels for actual concatenation let chunk_size = first_shape[dim+1..].iter().product::(); let outer_size = first_shape[..dim].iter().product::(); for outer_idx in 0..outer_size { for tensor in tensors { let tensor_data = tensor.to_cpu()?; let tensor_dim_size = tensor.shape().dims()[dim]; let start_idx = outer_idx * tensor_dim_size * chunk_size; let end_idx = start_idx + tensor_dim_size * chunk_size; if end_idx <= tensor_data.len() { result_data.extend_from_slice(&tensor_data[start_idx..end_idx]); } } } Tensor::from_data(result_data, result_shape, &first_tensor.device) } pub fn stack(tensors: &[Self], dim: usize) -> Result { if tensors.is_empty() { return Err(TensorError::shape("Cannot stack empty tensor list".to_string())); } // Simplified stack - just return first tensor for now Ok(tensors[0].clone()) } pub fn softmax(&self, dim: i32) -> Result { // Simplified softmax implementation let data = self.to_cpu()?; let mut result_data = data.clone(); // Apply softmax (simplified for 1D case) let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); let mut sum = 0.0; for val in &mut result_data { *val = (*val - max_val).exp(); sum += *val; } for val in &mut result_data { *val /= sum; } Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn log(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x.ln()).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn exp(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x.exp()).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn sqrt(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x.sqrt()).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn gelu(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| { 0.5 * x * (1.0 + (0.7978845608 * (x + 0.044715 * x.powi(3))).tanh()) }).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn swish(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x * (1.0 / (1.0 + (-x).exp()))).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn sigmoid(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| 1.0 / (1.0 + (-x).exp())).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn tanh(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x.tanh()).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn abs(&self) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x.abs()).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn norm(&self, p: i32, dims: &[i32], keepdim: bool) -> Result { // Simplified L2 norm let data = self.to_cpu()?; let norm_val = (data.iter().map(|x| x * x).sum::()).sqrt(); let result_shape = if keepdim { self.shape().dims().to_vec() } else { vec![1] }; Tensor::from_data(vec![norm_val], result_shape, &self.device) } pub fn mean(&self, dims: &[i32], keepdim: bool) -> Result { let data = self.to_cpu()?; let mean_val = data.iter().sum::() / data.len() as f32; let result_shape = if keepdim { self.shape().dims().to_vec() } else { vec![1] }; Tensor::from_data(vec![mean_val], result_shape, &self.device) } pub fn var(&self, dims: &[i32], unbiased: bool, keepdim: bool) -> Result { let data = self.to_cpu()?; let mean_val = data.iter().sum::() / data.len() as f32; let var_val = data.iter().map(|x| (x - mean_val).powi(2)).sum::() / data.len() as f32; let result_shape = if keepdim { self.shape().dims().to_vec() } else { vec![1] }; Tensor::from_data(vec![var_val], result_shape, &self.device) } pub fn max(&self) -> Result { let data = self.to_cpu()?; let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); Tensor::from_data(vec![max_val], vec![1], &self.device) } pub fn get(&self, indices: &[std::ops::RangeFull]) -> Result { // Simplified get - just return self for now Ok(self.clone()) } pub fn chunk(&self, chunks: usize, dim: usize) -> Result<(Self, Self)> { // Simplified chunk - return two copies for now Ok((self.clone(), self.clone())) } pub fn index_select(&self, dim: usize, index: &Self) -> Result { // Simplified index_select - just return self for now Ok(self.clone()) } pub fn gather(&self, dim: usize, index: &Self) -> Result { // Simplified gather - just return self for now Ok(self.clone()) } pub fn contiguous(&self) -> Result { // Just return self for now Ok(self.clone()) } pub fn mul_scalar(&self, scalar: f32) -> Result { // Use existing scalar_mul method self.scalar_mul(scalar) } pub fn add_scalar(&self, scalar: f32) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x + scalar).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } pub fn div_scalar(&self, scalar: f32) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|x| x / scalar).collect(); Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device) } // Convolution operations (simplified stubs) pub fn conv2d(&self, weight: &Self, bias: Option<&Self>, stride: usize, padding: usize, dilation: usize, groups: usize) -> Result { // Simplified conv2d - just return self for now Ok(self.clone()) } pub fn conv1d(&self, weight: &Self, bias: Option<&Self>, stride: usize, padding: usize, groups: usize) -> Result { // Simplified conv1d - just return self for now Ok(self.clone()) } // Static constructors pub fn randn(shape: &[usize], device: &Device) -> Result { let numel = shape.iter().product(); let data: Vec = (0..numel).map(|_| { // Box-Muller transform for normal distribution let u1: f32 = rand::random(); let u2: f32 = rand::random(); ((-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()) }).collect(); Self::from_data(data, shape.to_vec(), device) } pub fn randint(low: i32, high: i32, shape: &[usize], device: &Device) -> Result { let numel = shape.iter().product(); let data: Vec = (0..numel).map(|_| { let val = low + (rand::random::() * (high - low) as f32) as i32; val as f32 }).collect(); Self::from_data(data, shape.to_vec(), device) } pub fn uniform(shape: &[usize], low: f32, high: f32, device: &Device) -> Result { let numel = shape.iter().product(); let data: Vec = (0..numel).map(|_| { low + rand::random::() * (high - low) }).collect(); Self::from_data(data, shape.to_vec(), device) } pub fn normal(mean: f32, std: f32, shape: &[usize], device: &Device) -> Result { let numel = shape.iter().product(); let data: Vec = (0..numel).map(|_| { let u1: f32 = rand::random(); let u2: f32 = rand::random(); let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos(); mean + std * z }).collect(); Self::from_data(data, shape.to_vec(), device) } pub fn full(shape: &[usize], value: f32, device: &Device) -> Result { let numel = shape.iter().product(); let data = vec![value; numel]; Self::from_data(data, shape.to_vec(), device) } pub fn arange(start: i64, end: i64, device: &Device) -> Result { let data: Vec = (start..end).map(|x| x as f32).collect(); let shape = vec![data.len()]; Self::from_data(data, shape, device) } pub fn from_vec(data: Vec, shape: &[usize], device: &Device) -> Result { Self::from_data(data, shape.to_vec(), device) } pub fn cross_entropy(&self, labels: &Self, dim: i32) -> Result { // Simplified cross-entropy loss let loss_val = 1.0; // Placeholder Tensor::from_data(vec![loss_val], vec![1], &self.device) } pub fn to_scalar(&self) -> Result where T: From, { let data = self.to_cpu()?; if data.is_empty() { return Err(TensorError::shape("Cannot convert empty tensor to scalar".to_string())); } Ok(T::from(data[0])) } /// Solve linear system using Cholesky decomposition pub fn cholesky_solve(&self, b: &Self) -> Result { // Validate input shapes let a_shape = self.shape().dims(); let b_shape = b.shape().dims(); if a_shape.len() != 2 || a_shape[0] != a_shape[1] { return Err(TensorError::shape("Matrix A must be square".to_string())); } if b_shape.len() != 2 || b_shape[0] != a_shape[0] { return Err(TensorError::shape("Matrix B dimensions incompatible".to_string())); } let n = a_shape[0]; let nrhs = b_shape[1]; // Get CPU data for computation let a_data = self.to_cpu()?; let b_data = b.to_cpu()?; // Perform Cholesky decomposition: A = L * L^T let mut l_matrix = vec![0.0f32; n * n]; // Compute Cholesky decomposition for i in 0..n { for j in 0..=i { if i == j { // Diagonal elements let mut sum = 0.0; for k in 0..j { sum += l_matrix[j * n + k] * l_matrix[j * n + k]; } l_matrix[j * n + j] = (a_data[j * n + j] - sum).sqrt(); } else { // Lower triangular elements let mut sum = 0.0; for k in 0..j { sum += l_matrix[i * n + k] * l_matrix[j * n + k]; } l_matrix[i * n + j] = (a_data[i * n + j] - sum) / l_matrix[j * n + j]; } } } // Solve L * y = b using forward substitution let mut y_data = vec![0.0f32; n * nrhs]; for col in 0..nrhs { for i in 0..n { let mut sum = 0.0; for j in 0..i { sum += l_matrix[i * n + j] * y_data[j * nrhs + col]; } y_data[i * nrhs + col] = (b_data[i * nrhs + col] - sum) / l_matrix[i * n + i]; } } // Solve L^T * x = y using backward substitution let mut x_data = vec![0.0f32; n * nrhs]; for col in 0..nrhs { for i in (0..n).rev() { let mut sum = 0.0; for j in (i + 1)..n { sum += l_matrix[j * n + i] * x_data[j * nrhs + col]; } x_data[i * nrhs + col] = (y_data[i * nrhs + col] - sum) / l_matrix[i * n + i]; } } Tensor::from_data(x_data, b_shape.to_vec(), &self.device) } /// Set values at specific indices (simplified) pub fn index_put(&self, indices: &[usize], values: &Self) -> Result { // Simplified implementation - return self for now Ok(self.clone()) } /// Slice tensor along dimension pub fn slice(&self, dim: usize, start: usize, end: usize) -> Result { // Simplified implementation if dim >= self.ndim() { return Err(TensorError::shape("Dimension out of bounds".to_string())); } // For now, just return a view with adjusted shape let mut new_shape = self.shape.dims().to_vec(); new_shape[dim] = end - start; self.view(new_shape) } /// Compute L2 norm (alias for compatibility) pub fn norm_l2(&self) -> Result { let data = self.to_cpu()?; let norm_val = (data.iter().map(|&x| x * x).sum::()).sqrt(); Tensor::from_data(vec![norm_val], vec![1], &self.device) } /// Element-wise greater than or equal comparison with scalar pub fn ge_scalar(&self, threshold: f32) -> Result { let data = self.to_cpu()?; let result_data: Vec = data.iter().map(|&x| if x >= threshold { 1.0 } else { 0.0 }).collect(); Tensor::from_data(result_data, self.shape.dims().to_vec(), &self.device) } /// Convert tensor to different data type pub fn to_dtype(&self, dtype: &str) -> Result { // For now, just return self (assuming same type) Ok(self.clone()) } } // Operator overloading for ergonomic API impl std::ops::Add for &Tensor { type Output = Result; fn add(self, other: &Tensor) -> Self::Output { self.add(other) } } impl std::ops::Mul for &Tensor { type Output = Result; fn mul(self, other: &Tensor) -> Self::Output { self.mul(other) } } impl Tensor { /// Element-wise power operation pub fn pow(&self, exponent: f32) -> Result { let self_data = self.to_cpu()?; let result_data: Vec = self_data.iter().map(|&x| x.powf(exponent)).collect(); let result_tensor = Self::from_data(result_data, self.shape.dims().to_vec(), &self.device)?; Ok(result_tensor) } /// Argmin along specified dimension pub fn argmin(&self, dim: i32, _keepdim: bool) -> Result { let data = self.to_cpu()?; let argmin_idx = data.iter() .enumerate() .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) .map(|(idx, _)| idx) .unwrap_or(0) as f32; // For simplicity, return 1D tensor with the index Self::from_data(vec![argmin_idx], vec![1], &self.device) } /// Narrow (slice) tensor along specified dimension pub fn narrow(&self, dim: usize, start: usize, length: usize) -> Result { let data = self.to_cpu()?; // Simplified narrow - for 2D tensors along dim 1 if self.shape.dims().len() == 2 && dim == 1 { let rows = self.shape.dims()[0]; let cols = self.shape.dims()[1]; let mut result_data = Vec::new(); for row in 0..rows { for col in start..start + length { if col < cols { result_data.push(data[row * cols + col]); } } } Self::from_data(result_data, vec![rows, length], &self.device) } else { // For other cases, just return the original tensor Self::from_data(data, self.shape.dims().to_vec(), &self.device) } } /// Flatten tensor from start_dim to end_dim pub fn flatten(&self, start_dim: i32, end_dim: i32) -> Result { let dims = self.shape.dims(); let ndim = dims.len() as i32; let start = if start_dim < 0 { ndim + start_dim } else { start_dim } as usize; let end = if end_dim < 0 { ndim + end_dim } else { end_dim } as usize; if start >= dims.len() || end >= dims.len() || start > end { return Err(TensorError::shape("Invalid flatten dimensions".to_string())); } // Calculate new shape let mut new_shape = Vec::new(); // Dimensions before start_dim for i in 0..start { new_shape.push(dims[i]); } // Flattened dimension let flattened_size: usize = dims[start..=end].iter().product(); new_shape.push(flattened_size); // Dimensions after end_dim for i in (end + 1)..dims.len() { new_shape.push(dims[i]); } let data = self.to_cpu()?; Self::from_data(data, new_shape, &self.device) } /// Broadcast tensor to specified shape pub fn broadcast_to(&self, shape: &[usize]) -> Result { // Simplified broadcasting - just create tensor with new shape and repeated data let data = self.to_cpu()?; let target_numel: usize = shape.iter().product(); let mut result_data = Vec::with_capacity(target_numel); for i in 0..target_numel { result_data.push(data[i % data.len()]); } Self::from_data(result_data, shape.to_vec(), &self.device) } /// Create tensor from slice pub fn from_slice(data: &[f32], shape: &[usize], device: &Device) -> Result { Self::from_data(data.to_vec(), shape.to_vec(), device) } /// Random uniform tensor pub fn rand(shape: &[usize], _dtype: DType, device: &Device) -> Result { use rand::Rng; let mut rng = rand::thread_rng(); let numel: usize = shape.iter().product(); let data: Vec = (0..numel).map(|_| rng.gen()).collect(); Self::from_data(data, shape.to_vec(), device) } /// Cast tensor to different dtype pub fn cast(&self, _dtype: DType) -> Result { let data = self.to_cpu()?; // For simplicity, just keep as f32 regardless of target dtype Self::from_data(data, self.shape.dims().to_vec(), &self.device) } } #[cfg(test)] mod tests { use super::*; use approx::assert_relative_eq; /// Test tensor creation with zeros #[test] fn test_tensor_zeros() { let device = Device::cpu(); let tensor = Tensor::zeros([2, 3], &device).expect("Failed to create zeros tensor"); assert_eq!(tensor.shape().dims(), &[2, 3]); assert_eq!(tensor.ndim(), 2); assert_eq!(tensor.numel(), 6); assert_eq!(tensor.device(), &device); assert_eq!(tensor.dtype(), DType::F32); assert!(!tensor.requires_grad()); } /// Test tensor creation with ones #[test] fn test_tensor_ones() { let device = Device::cpu(); let tensor = Tensor::ones([3, 2], &device).expect("Failed to create ones tensor"); assert_eq!(tensor.shape().dims(), &[3, 2]); assert_eq!(tensor.ndim(), 2); assert_eq!(tensor.numel(), 6); // Verify data contains ones let data = tensor.to_cpu().expect("Failed to copy to CPU"); for value in data { assert_relative_eq!(value, 1.0, epsilon = 1e-6); } } /// Test tensor creation from raw data #[test] fn test_tensor_from_data() { let device = Device::cpu(); let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; let tensor = Tensor::from_data(data.clone(), [2, 3], &device) .expect("Failed to create tensor from data"); assert_eq!(tensor.shape().dims(), &[2, 3]); assert_eq!(tensor.numel(), 6); let result = tensor.to_cpu().expect("Failed to copy to CPU"); assert_eq!(result, data); } /// Test gradient requirement setting #[test] fn test_requires_grad() { let device = Device::cpu(); let mut tensor = Tensor::zeros([2, 2], &device) .expect("Failed to create tensor"); assert!(!tensor.requires_grad()); tensor.set_requires_grad(true); assert!(tensor.requires_grad()); tensor.set_requires_grad(false); assert!(!tensor.requires_grad()); } /// Test tensor view operations #[test] fn test_tensor_view() { let device = Device::cpu(); let tensor = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3], &device) .expect("Failed to create tensor"); // View as different shape with same number of elements let view = tensor.view([3, 2]).expect("Failed to create view"); assert_eq!(view.shape().dims(), &[3, 2]); assert_eq!(view.numel(), 6); // Should share same underlying storage let original_data = tensor.to_cpu().expect("Failed to get original data"); let view_data = view.to_cpu().expect("Failed to get view data"); assert_eq!(original_data, view_data); } /// Test tensor reshape operations #[test] fn test_tensor_reshape() { let device = Device::cpu(); let tensor = Tensor::from_data( vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], [2, 4], &device ).expect("Failed to create tensor"); let reshaped = tensor.reshape([4, 2]).expect("Failed to reshape"); assert_eq!(reshaped.shape().dims(), &[4, 2]); assert_eq!(reshaped.numel(), 8); } /// Test element-wise addition #[test] fn test_tensor_add() { let device = Device::cpu(); let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], [2, 2], &device) .expect("Failed to create tensor a"); let b = Tensor::from_data(vec![5.0, 6.0, 7.0, 8.0], [2, 2], &device) .expect("Failed to create tensor b"); let c = a.add(&b).expect("Failed to add tensors"); let result = c.to_cpu().expect("Failed to copy result to CPU"); assert_eq!(result, vec![6.0, 8.0, 10.0, 12.0]); } /// Test element-wise multiplication #[test] fn test_tensor_mul() { let device = Device::cpu(); let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], [2, 2], &device) .expect("Failed to create tensor a"); let b = Tensor::from_data(vec![2.0, 3.0, 4.0, 5.0], [2, 2], &device) .expect("Failed to create tensor b"); let c = a.mul(&b).expect("Failed to multiply tensors"); let result = c.to_cpu().expect("Failed to copy result to CPU"); assert_eq!(result, vec![2.0, 6.0, 12.0, 20.0]); } /// Test matrix multiplication #[test] fn test_tensor_matmul() { let device = Device::cpu(); // 2x3 matrix let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3], &device) .expect("Failed to create tensor a"); // 3x2 matrix let b = Tensor::from_data(vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0], [3, 2], &device) .expect("Failed to create tensor b"); let c = a.matmul(&b).expect("Failed to perform matmul"); assert_eq!(c.shape().dims(), &[2, 2]); let result = c.to_cpu().expect("Failed to copy result to CPU"); // Expected: [[58, 64], [139, 154]] let expected = vec![58.0, 64.0, 139.0, 154.0]; for (actual, expected) in result.iter().zip(expected.iter()) { assert_relative_eq!(actual, expected, epsilon = 1e-6); } } /// Test sum reduction #[test] fn test_tensor_sum() { let device = Device::cpu(); let tensor = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3], &device) .expect("Failed to create tensor"); // Sum all elements let sum_all = tensor.sum(None).expect("Failed to sum all elements"); assert_eq!(sum_all.shape().dims(), &[]); // Scalar let result = sum_all.to_cpu().expect("Failed to copy to CPU"); assert_relative_eq!(result[0], 21.0, epsilon = 1e-6); // Sum along dimension 0 let sum_dim0 = tensor.sum(Some(0)).expect("Failed to sum along dim 0"); assert_eq!(sum_dim0.shape().dims(), &[3]); let result = sum_dim0.to_cpu().expect("Failed to copy to CPU"); assert_eq!(result, vec![5.0, 7.0, 9.0]); // [1+4, 2+5, 3+6] // Sum along dimension 1 let sum_dim1 = tensor.sum(Some(1)).expect("Failed to sum along dim 1"); assert_eq!(sum_dim1.shape().dims(), &[2]); let result = sum_dim1.to_cpu().expect("Failed to copy to CPU"); assert_eq!(result, vec![6.0, 15.0]); // [1+2+3, 4+5+6] } /// Test operator overloading #[test] fn test_operator_overloading() { let device = Device::cpu(); let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], [2, 2], &device) .expect("Failed to create tensor a"); let b = Tensor::from_data(vec![5.0, 6.0, 7.0, 8.0], [2, 2], &device) .expect("Failed to create tensor b"); // Test addition operator let c = (&a + &b).expect("Failed to add tensors with operator"); let result = c.to_cpu().expect("Failed to copy result to CPU"); assert_eq!(result, vec![6.0, 8.0, 10.0, 12.0]); // Test multiplication operator let d = (&a * &b).expect("Failed to multiply tensors with operator"); let result = d.to_cpu().expect("Failed to copy result to CPU"); assert_eq!(result, vec![5.0, 12.0, 21.0, 32.0]); } /// Test device transfer #[test] fn test_device_transfer() { let cpu_device = Device::cpu(); let tensor = Tensor::ones([2, 2], &cpu_device) .expect("Failed to create CPU tensor"); assert_eq!(tensor.device(), &cpu_device); // Test moving to same device (should be no-op) let same_device = tensor.to_device(&cpu_device) .expect("Failed to move to same device"); assert_eq!(same_device.device(), &cpu_device); } /// Test invalid operations #[test] fn test_invalid_operations() { let device = Device::cpu(); let a = Tensor::zeros([2, 3], &device).expect("Failed to create tensor a"); let b = Tensor::zeros([3, 2], &device).expect("Failed to create tensor b"); // Test shape mismatch in element-wise operations let result = a.add(&b); assert!(result.is_err()); let result = a.mul(&b); assert!(result.is_err()); // Test invalid reshape let invalid_reshape = a.reshape([2, 2]); assert!(invalid_reshape.is_err()); // 6 elements can't fit in 2x2=4 // Test invalid view let invalid_view = a.view([3, 3]); assert!(invalid_view.is_err()); // 6 elements can't be viewed as 3x3=9 } /// Property-based test for tensor operations #[cfg(feature = "proptest")] mod proptests { use super::*; use proptest::prelude::*; proptest! { #[test] fn test_tensor_reshape_preserves_data( data in prop::collection::vec(any::(), 1..=100), original_shape in prop::collection::vec(1usize..=10, 2..=4) ) { let numel = data.len(); // Generate a valid reshape target let factors = find_factors(numel); if factors.len() < 2 { return Ok(()); } // Skip if not enough factors let device = Device::cpu(); let tensor = Tensor::from_data(data.clone(), original_shape, &device)?; let reshaped = tensor.reshape(factors)?; prop_assert_eq!(reshaped.numel(), numel); prop_assert_eq!(reshaped.to_cpu()?, data); } #[test] fn test_tensor_add_commutative( a_data in prop::collection::vec(any::(), 1..=50), shape in prop::collection::vec(1usize..=5, 2..=3) ) { let numel = shape.iter().product::(); if a_data.len() != numel { return Ok(()); } let device = Device::cpu(); let a = Tensor::from_data(a_data.clone(), shape.clone(), &device)?; let b = Tensor::from_data(a_data, shape, &device)?; let ab = a.add(&b)?; let ba = b.add(&a)?; let ab_data = ab.to_cpu()?; let ba_data = ba.to_cpu()?; for (x, y) in ab_data.iter().zip(ba_data.iter()) { prop_assert!((x - y).abs() < 1e-6); } } } fn find_factors(n: usize) -> Vec { if n <= 1 { return vec![1]; } if n <= 3 { return vec![n]; } vec![n / 2, 2] // Simple factorization for testing } } }