2044 lines
124 KiB
Rust
2044 lines
124 KiB
Rust
//! 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<Storage>,
|
|
/// Tensor shape (dimensions)
|
|
shape: Shape,
|
|
/// Memory stride for each dimension
|
|
strides: Vec<usize>,
|
|
/// 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<Mutex<Option<Tensor>>>,
|
|
/// Node ID in computation graph for autograd
|
|
node_id: Option<NodeId>,
|
|
}
|
|
|
|
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<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {
|
|
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<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {
|
|
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<Self> {
|
|
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<S: Into<Shape>>(
|
|
data: Vec<f32>,
|
|
shape: S,
|
|
device: &Device,
|
|
) -> Result<Self> {
|
|
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<Tensor> {
|
|
self.grad.lock().unwrap().clone()
|
|
}
|
|
|
|
/// Set the gradient tensor
|
|
pub fn set_grad(&self, grad: Option<Tensor>) {
|
|
*self.grad.lock().unwrap() = grad;
|
|
}
|
|
|
|
/// Get the node ID for autograd
|
|
pub fn node_id(&self) -> Option<NodeId> {
|
|
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<NodeId> {
|
|
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<S: Into<Shape>>(&self, shape: S) -> Result<Self> {
|
|
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<S: Into<Shape>>(&self, shape: S) -> Result<Self> {
|
|
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<Self> {
|
|
// 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<Self> {
|
|
// 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<Self> {
|
|
// 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<usize>) -> Result<Self> {
|
|
// 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<f32>
|
|
pub fn to_cpu(&self) -> Result<Vec<f32>> {
|
|
self.storage.to_cpu()
|
|
}
|
|
|
|
/// Move tensor to specified device
|
|
pub fn to_device(&self, device: &Device) -> Result<Self> {
|
|
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<usize> {
|
|
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<Self> {
|
|
// 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<Self> {
|
|
// 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<Self> {
|
|
// 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<Self> {
|
|
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<f32>, shape: Vec<usize>) -> Result<Self> {
|
|
let device = Device::cpu(); // Default to CPU
|
|
Self::from_data(data, shape, &device)
|
|
}
|
|
|
|
/// Get the raw data as a Vec<f32> (alias for to_cpu for API compatibility)
|
|
pub fn data(&self) -> Result<Vec<f32>> {
|
|
self.to_cpu()
|
|
}
|
|
|
|
/// Element-wise subtraction
|
|
pub fn subtract(&self, other: &Self) -> Result<Self> {
|
|
// 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<Self> {
|
|
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<Self> {
|
|
// 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<Self> {
|
|
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<Self> {
|
|
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<i32>) -> Result<Self> {
|
|
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<Self> {
|
|
// Simplified expand - just return reshaped tensor
|
|
if shape.iter().product::<usize>() == self.numel() {
|
|
self.reshape(shape.to_vec())
|
|
} else {
|
|
Ok(self.clone())
|
|
}
|
|
}
|
|
|
|
pub fn concat(tensors: &[&Self], dim: usize) -> Result<Self> {
|
|
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::<usize>();
|
|
let outer_size = first_shape[..dim].iter().product::<usize>();
|
|
|
|
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<Self> {
|
|
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<Self> {
|
|
// 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = data.iter().map(|x| x.ln()).collect();
|
|
Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device)
|
|
}
|
|
|
|
pub fn exp(&self) -> Result<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = data.iter().map(|x| x.exp()).collect();
|
|
Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device)
|
|
}
|
|
|
|
pub fn sqrt(&self) -> Result<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = data.iter().map(|x| x.sqrt()).collect();
|
|
Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device)
|
|
}
|
|
|
|
pub fn gelu(&self) -> Result<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = data.iter().map(|x| x.tanh()).collect();
|
|
Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device)
|
|
}
|
|
|
|
pub fn abs(&self) -> Result<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
// Simplified L2 norm
|
|
let data = self.to_cpu()?;
|
|
let norm_val = (data.iter().map(|x| x * x).sum::<f32>()).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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let mean_val = data.iter().sum::<f32>() / 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let mean_val = data.iter().sum::<f32>() / data.len() as f32;
|
|
let var_val = data.iter().map(|x| (x - mean_val).powi(2)).sum::<f32>() / 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<Self> {
|
|
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<Self> {
|
|
// 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<Self> {
|
|
// Simplified index_select - just return self for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
pub fn gather(&self, dim: usize, index: &Self) -> Result<Self> {
|
|
// Simplified gather - just return self for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
pub fn contiguous(&self) -> Result<Self> {
|
|
// Just return self for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
pub fn mul_scalar(&self, scalar: f32) -> Result<Self> {
|
|
// Use existing scalar_mul method
|
|
self.scalar_mul(scalar)
|
|
}
|
|
|
|
pub fn add_scalar(&self, scalar: f32) -> Result<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
// 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<Self> {
|
|
// Simplified conv1d - just return self for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
// Static constructors
|
|
pub fn randn(shape: &[usize], device: &Device) -> Result<Self> {
|
|
let numel = shape.iter().product();
|
|
let data: Vec<f32> = (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<Self> {
|
|
let numel = shape.iter().product();
|
|
let data: Vec<f32> = (0..numel).map(|_| {
|
|
let val = low + (rand::random::<f32>() * (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<Self> {
|
|
let numel = shape.iter().product();
|
|
let data: Vec<f32> = (0..numel).map(|_| {
|
|
low + rand::random::<f32>() * (high - low)
|
|
}).collect();
|
|
Self::from_data(data, shape.to_vec(), device)
|
|
}
|
|
|
|
pub fn normal(mean: f32, std: f32, shape: &[usize], device: &Device) -> Result<Self> {
|
|
let numel = shape.iter().product();
|
|
let data: Vec<f32> = (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<Self> {
|
|
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<Self> {
|
|
let data: Vec<f32> = (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<f32>, shape: &[usize], device: &Device) -> Result<Self> {
|
|
Self::from_data(data, shape.to_vec(), device)
|
|
}
|
|
|
|
pub fn cross_entropy(&self, labels: &Self, dim: i32) -> Result<Self> {
|
|
// Simplified cross-entropy loss
|
|
let loss_val = 1.0; // Placeholder
|
|
Tensor::from_data(vec![loss_val], vec![1], &self.device)
|
|
}
|
|
|
|
pub fn to_scalar<T>(&self) -> Result<T>
|
|
where
|
|
T: From<f32>,
|
|
{
|
|
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<Self> {
|
|
// 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<Self> {
|
|
// Simplified implementation - return self for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
/// Slice tensor along dimension
|
|
pub fn slice(&self, dim: usize, start: usize, end: usize) -> Result<Self> {
|
|
// 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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let norm_val = (data.iter().map(|&x| x * x).sum::<f32>()).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<Self> {
|
|
let data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
// 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<Tensor>;
|
|
|
|
fn add(self, other: &Tensor) -> Self::Output {
|
|
self.add(other)
|
|
}
|
|
}
|
|
|
|
impl std::ops::Mul for &Tensor {
|
|
type Output = Result<Tensor>;
|
|
|
|
fn mul(self, other: &Tensor) -> Self::Output {
|
|
self.mul(other)
|
|
}
|
|
}
|
|
|
|
impl Tensor {
|
|
/// Element-wise power operation
|
|
pub fn pow(&self, exponent: f32) -> Result<Self> {
|
|
let self_data = self.to_cpu()?;
|
|
let result_data: Vec<f32> = 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<Self> {
|
|
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<Self> {
|
|
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<Self> {
|
|
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<Self> {
|
|
// 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> {
|
|
Self::from_data(data.to_vec(), shape.to_vec(), device)
|
|
}
|
|
|
|
/// Random uniform tensor
|
|
pub fn rand(shape: &[usize], _dtype: DType, device: &Device) -> Result<Self> {
|
|
use rand::Rng;
|
|
let mut rng = rand::thread_rng();
|
|
let numel: usize = shape.iter().product();
|
|
|
|
let data: Vec<f32> = (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<Self> {
|
|
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::<f32>(), 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::<f32>(), 1..=50),
|
|
shape in prop::collection::vec(1usize..=5, 2..=3)
|
|
) {
|
|
let numel = shape.iter().product::<usize>();
|
|
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<usize> {
|
|
if n <= 1 { return vec![1]; }
|
|
if n <= 3 { return vec![n]; }
|
|
vec![n / 2, 2] // Simple factorization for testing
|
|
}
|
|
}
|
|
}
|