2 lines
25 KiB
Rust
2 lines
25 KiB
Rust
//! Vision Tensor Implementation\n//!\n//! Production-grade tensor operations for computer vision models,\n//! integrating with the rtx-tensor ecosystem for real GPU acceleration.\n\nuse serde::{Deserialize, Serialize};\nuse std::fmt;\nuse std::sync::Arc;\nuse tracing::{debug, warn, error, info};\n\n/// Re-export from rtx-tensor for core functionality\nuse rtx_tensor::{\n Tensor as CoreTensor, Device as CoreDevice, Shape as CoreShape, DType as CoreDType,\n TensorError as CoreTensorError, Result as CoreResult\n};\n\n/// Vision tensor error types\n#[derive(Debug, thiserror::Error)]\npub enum VisionTensorError {\n #[error(\"Shape mismatch: {0}\")]\n ShapeMismatch(String),\n #[error(\"Device mismatch: expected {expected:?}, got {actual:?}\")]\n DeviceMismatch { expected: Device, actual: Device },\n #[error(\"Invalid operation: {0}\")]\n InvalidOperation(String),\n #[error(\"Out of memory: requested {size} bytes\")]\n OutOfMemory { size: usize },\n #[error(\"Tensor computation failed: {0}\")]\n ComputationError(String),\n #[error(\"Invalid tensor index: {0}\")]\n InvalidIndex(String),\n #[error(\"Unsupported data type: {0}\")]\n UnsupportedDType(String),\n #[error(\"GPU error: {0}\")]\n GpuError(String),\n #[error(\"Core tensor error: {0}\")]\n CoreError(#[from] CoreTensorError),\n}\n\npub type Result<T> = std::result::Result<T, VisionTensorError>;\n\n/// Device wrapper for vision operations\n#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]\npub enum Device {\n /// CPU execution\n Cpu,\n /// CUDA GPU with device ID\n Cuda(i32),\n}\n\nimpl Device {\n pub fn cpu() -> Self {\n Self::Cpu\n }\n\n pub fn cuda(device_id: i32) -> Self {\n Self::Cuda(device_id)\n }\n\n pub fn is_cuda(&self) -> bool {\n matches!(self, Self::Cuda(_))\n }\n\n pub fn device_id(&self) -> Option<i32> {\n match self {\n Self::Cpu => None,\n Self::Cuda(id) => Some(*id),\n }\n }\n\n /// Check if device is available\n pub fn is_available(&self) -> bool {\n match self {\n Self::Cpu => true,\n Self::Cuda(device_id) => {\n // In production, would check actual GPU availability\n *device_id >= 0 && *device_id < 8 // Assume max 8 GPUs\n }\n }\n }\n\n /// Convert to core device\n fn to_core_device(&self) -> CoreDevice {\n match self {\n Self::Cpu => CoreDevice::cpu(),\n Self::Cuda(id) => CoreDevice::cuda(*id),\n }\n }\n\n /// Convert from core device\n fn from_core_device(core_device: &CoreDevice) -> Self {\n match core_device {\n CoreDevice::cuda(0).unwrap_or(Device::default()) => Self::Cpu,\n CoreDevice::Cuda(id) => Self::Cuda(*id),\n }\n }\n}\n\nimpl fmt::Display for Device {\n fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {\n match self {\n Self::Cpu => write!(f, \"cpu\"),\n Self::Cuda(id) => write!(f, \"cuda:{}\", id),\n }\n }\n}\n\n/// Shape wrapper for vision tensors\n#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]\npub struct Shape {\n dims: Vec<usize>,\n}\n\nimpl Shape {\n pub fn new(dims: Vec<usize>) -> Self {\n Self { dims }\n }\n\n pub fn from_slice(dims: &[usize]) -> Self {\n Self::new(dims.to_vec())\n }\n\n pub fn dims(&self) -> &[usize] {\n &self.dims\n }\n\n pub fn ndim(&self) -> usize {\n self.dims.len()\n }\n\n pub fn numel(&self) -> usize {\n self.dims.iter().product()\n }\n\n pub fn is_empty(&self) -> bool {\n self.dims.is_empty() || self.dims.iter().any(|&d| d == 0)\n }\n\n pub fn is_scalar(&self) -> bool {\n self.dims.is_empty()\n }\n\n pub fn is_vector(&self) -> bool {\n self.dims.len() == 1\n }\n\n pub fn is_matrix(&self) -> bool {\n self.dims.len() == 2\n }\n\n /// Check if this is a valid image tensor shape [N, C, H, W] or [C, H, W]\n pub fn is_image_like(&self) -> bool {\n matches!(self.dims.len(), 3 | 4) && self.dims.iter().all(|&d| d > 0)\n }\n\n /// Get batch size for 4D tensors [N, C, H, W]\n pub fn batch_size(&self) -> Option<usize> {\n if self.dims.len() == 4 {\n Some(self.dims[0])\n } else {\n None\n }\n }\n\n /// Get number of channels for 3D/4D tensors\n pub fn channels(&self) -> Option<usize> {\n match self.dims.len() {\n 3 => Some(self.dims[0]), // [C, H, W]\n 4 => Some(self.dims[1]), // [N, C, H, W]\n _ => None,\n }\n }\n\n /// Get height for 3D/4D tensors\n pub fn height(&self) -> Option<usize> {\n match self.dims.len() {\n 3 => Some(self.dims[1]), // [C, H, W]\n 4 => Some(self.dims[2]), // [N, C, H, W]\n _ => None,\n }\n }\n\n /// Get width for 3D/4D tensors\n pub fn width(&self) -> Option<usize> {\n match self.dims.len() {\n 3 => Some(self.dims[2]), // [C, H, W]\n 4 => Some(self.dims[3]), // [N, C, H, W]\n _ => None,\n }\n }\n\n /// Validate shape for vision operations\n pub fn validate_vision_shape(&self) -> Result<()> {\n if self.is_empty() {\n return Err(VisionTensorError::InvalidOperation(\n \"Empty tensor shapes not supported for vision operations\".to_string()\n ));\n }\n\n if self.dims.len() < 3 || self.dims.len() > 4 {\n return Err(VisionTensorError::InvalidOperation(\n \"Vision tensors must have 3 or 4 dimensions [C, H, W] or [N, C, H, W]\".to_string()\n ));\n }\n\n if let Some(h) = self.height() {\n if h < 1 || h > 10000 {\n return Err(VisionTensorError::InvalidOperation(\n format!(\"Invalid height: {}. Must be between 1 and 10000\", h)\n ));\n }\n }\n\n if let Some(w) = self.width() {\n if w < 1 || w > 10000 {\n return Err(VisionTensorError::InvalidOperation(\n format!(\"Invalid width: {}. Must be between 1 and 10000\", w)\n ));\n }\n }\n\n Ok(())\n }\n\n /// Convert to core shape\n fn to_core_shape(&self) -> CoreShape {\n CoreShape::from(self.dims.clone())\n }\n\n /// Convert from core shape\n fn from_core_shape(core_shape: &CoreShape) -> Self {\n Self::new(core_shape.dims().to_vec())\n }\n}\n\n// Implement conversions for Shape\nimpl From<Vec<usize>> for Shape {\n fn from(dims: Vec<usize>) -> Self {\n Self::new(dims)\n }\n}\n\nimpl From<&[usize]> for Shape {\n fn from(dims: &[usize]) -> Self {\n Self::from_slice(dims)\n }\n}\n\nimpl<const N: usize> From<[usize; N]> for Shape {\n fn from(dims: [usize; N]) -> Self {\n Self::new(dims.to_vec())\n }\n}\n\nimpl<const N: usize> From<&[usize; N]> for Shape {\n fn from(dims: &[usize; N]) -> Self {\n Self::new(dims.to_vec())\n }\n}\n\n/// Data types for vision tensors\n#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]\npub enum DType {\n /// 32-bit float (most common for vision models)\n F32,\n /// 64-bit float (high precision)\n F64,\n /// 16-bit float (memory efficient)\n F16,\n /// 32-bit signed integer\n I32,\n /// 64-bit signed integer\n I64,\n /// 8-bit unsigned integer (for images)\n U8,\n /// Boolean\n Bool,\n}\n\nimpl DType {\n /// Get size in bytes\n pub fn size(&self) -> usize {\n match self {\n Self::F32 | Self::I32 => 4,\n Self::F64 | Self::I64 => 8,\n Self::F16 => 2,\n Self::U8 | Self::Bool => 1,\n }\n }\n\n /// Check if this is a floating point type\n pub fn is_float(&self) -> bool {\n matches!(self, Self::F32 | Self::F64 | Self::F16)\n }\n\n /// Check if this is an integer type\n pub fn is_integer(&self) -> bool {\n matches!(self, Self::I32 | Self::I64 | Self::U8)\n }\n\n /// Convert to core dtype\n fn to_core_dtype(&self) -> CoreDType {\n match self {\n Self::F32 => CoreDType::F32,\n Self::F64 => CoreDType::F64,\n Self::F16 => CoreDType::F16,\n Self::I32 => CoreDType::I32,\n Self::I64 => CoreDType::I64,\n Self::U8 => CoreDType::U8,\n Self::Bool => CoreDType::Bool,\n }\n }\n\n /// Convert from core dtype\n fn from_core_dtype(core_dtype: &CoreDType) -> Self {\n match core_dtype {\n CoreDType::F32 => Self::F32,\n CoreDType::F64 => Self::F64,\n CoreDType::F16 => Self::F16,\n CoreDType::I32 => Self::I32,\n CoreDType::I64 => Self::I64,\n CoreDType::U8 => Self::U8,\n CoreDType::Bool => Self::Bool,\n }\n }\n}\n\n/// Vision tensor wrapper around rtx-tensor\n#[derive(Debug, Clone)]\npub struct Tensor {\n core_tensor: CoreTensor,\n shape: Shape,\n device: Device,\n dtype: DType,\n}\n\nimpl Tensor {\n /// Create tensor from core tensor\n fn from_core_tensor(core_tensor: CoreTensor) -> Result<Self> {\n let core_shape = core_tensor.shape();\n let core_device = core_tensor.device();\n let core_dtype = core_tensor.dtype();\n\n Ok(Self {\n core_tensor,\n shape: Shape::from_core_shape(core_shape),\n device: Device::from_core_device(core_device),\n dtype: DType::from_core_dtype(core_dtype),\n })\n }\n\n /// Create zero tensor\n pub fn zeros<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {\n let shape = shape.into();\n shape.validate_vision_shape()?;\n\n let core_tensor = CoreTensor::zeros(shape.dims(), &device.to_core_device())\n .map_err(VisionTensorError::CoreError)?;\n\n debug!(\"Created zeros tensor with shape {:?} on {}\", shape.dims(), device);\n Self::from_core_tensor(core_tensor)\n }\n\n /// Create ones tensor\n pub fn ones<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {\n let shape = shape.into();\n shape.validate_vision_shape()?;\n\n let core_tensor = CoreTensor::ones(shape.dims(), &device.to_core_device())\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Create random tensor with normal distribution\n pub fn randn<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {\n let shape = shape.into();\n shape.validate_vision_shape()?;\n\n let core_tensor = CoreTensor::randn(shape.dims(), &device.to_core_device())\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Create tensor from raw data\n pub fn from_data(data: Vec<f32>, shape: Vec<usize>, device: &Device) -> Result<Self> {\n let shape_obj = Shape::from(shape.clone());\n shape_obj.validate_vision_shape()?;\n\n if data.len() != shape_obj.numel() {\n return Err(VisionTensorError::ShapeMismatch(\n format!(\"Data length {} doesn't match shape {:?} (expected {})\",\n data.len(), shape_obj.dims(), shape_obj.numel())\n ));\n }\n\n let core_tensor = CoreTensor::from_slice(&data, &shape, &device.to_core_device())\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Create tensor from image data (H, W, C -> C, H, W conversion)\n pub fn from_image(image_data: &[u8], height: usize, width: usize, channels: usize, device: &Device) -> Result<Self> {\n if image_data.len() != height * width * channels {\n return Err(VisionTensorError::ShapeMismatch(\n format!(\"Image data length {} doesn't match dimensions {}x{}x{}\",\n image_data.len(), height, width, channels)\n ));\n }\n\n // Convert from HWC to CHW format and normalize to [0, 1]\n let mut tensor_data = vec![0.0f32; channels * height * width];\n for c in 0..channels {\n for h in 0..height {\n for w in 0..width {\n let src_idx = h * width * channels + w * channels + c; // HWC\n let dst_idx = c * height * width + h * width + w; // CHW\n tensor_data[dst_idx] = image_data[src_idx] as f32 / 255.0;\n }\n }\n }\n\n Self::from_data(tensor_data, vec![channels, height, width], device)\n }\n\n /// Get tensor shape\n pub fn shape(&self) -> &Shape {\n &self.shape\n }\n\n /// Get tensor device\n pub fn device(&self) -> &Device {\n &self.device\n }\n\n /// Get data type\n pub fn dtype(&self) -> DType {\n self.dtype\n }\n\n /// Check if requires gradient\n pub fn requires_grad(&self) -> bool {\n self.core_tensor.requires_grad()\n }\n\n /// Set gradient requirement\n pub fn set_requires_grad(&mut self, requires_grad: bool) {\n self.core_tensor.set_requires_grad(requires_grad);\n }\n\n /// Get number of elements\n pub fn numel(&self) -> usize {\n self.shape.numel()\n }\n\n /// Get data as CPU Vec\n pub fn to_vec(&self) -> Result<Vec<f32>> {\n self.core_tensor.data()\n .iter()\n .map(|&x| Ok(x))\n .collect::<Result<Vec<f32>>>()\n .map_err(|_| VisionTensorError::ComputationError(\"Failed to convert to vec\".to_string()))\n }\n\n /// Move tensor to different device\n pub fn to_device(&self, device: &Device) -> Result<Self> {\n if self.device == *device {\n return Ok(self.clone());\n }\n\n if !device.is_available() {\n return Err(VisionTensorError::GpuError(\n format!(\"Device {} is not available\", device)\n ));\n }\n\n let core_tensor = self.core_tensor.to_device(&device.to_core_device())\n .map_err(VisionTensorError::CoreError)?;\n\n info!(\"Moved tensor from {} to {}\", self.device, device);\n Self::from_core_tensor(core_tensor)\n }\n\n /// Reshape tensor\n pub fn reshape<S: AsRef<[usize]>>(&self, new_shape: S) -> Result<Self> {\n let new_shape = Shape::from_slice(new_shape.as_ref());\n\n if self.shape.numel() != new_shape.numel() {\n return Err(VisionTensorError::ShapeMismatch(\n format!(\"Cannot reshape tensor of size {} to size {}\",\n self.shape.numel(), new_shape.numel())\n ));\n }\n\n let core_tensor = self.core_tensor.reshape(new_shape.dims())\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Transpose tensor along specified dimensions\n pub fn transpose(&self, dim0: i32, dim1: i32) -> Result<Self> {\n let rank = self.shape.ndim() as i32;\n\n // Handle negative dimensions\n let dim0 = if dim0 < 0 { rank + dim0 } else { dim0 };\n let dim1 = if dim1 < 0 { rank + dim1 } else { dim1 };\n\n if dim0 < 0 || dim0 >= rank || dim1 < 0 || dim1 >= rank {\n return Err(VisionTensorError::InvalidOperation(\n \"Invalid dimensions for transpose\".to_string()\n ));\n }\n\n let core_tensor = self.core_tensor.transpose()\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Element-wise addition\n pub fn add(&self, other: &Self) -> Result<Self> {\n self.check_compatible_shape(other)?;\n self.check_compatible_device(other)?;\n\n let core_tensor = self.core_tensor.add(&other.core_tensor)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Element-wise subtraction\n pub fn sub(&self, other: &Self) -> Result<Self> {\n self.check_compatible_shape(other)?;\n self.check_compatible_device(other)?;\n\n let core_tensor = self.core_tensor.sub(&other.core_tensor)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Element-wise multiplication\n pub fn mul(&self, other: &Self) -> Result<Self> {\n self.check_compatible_shape(other)?;\n self.check_compatible_device(other)?;\n\n let core_tensor = self.core_tensor.mul(&other.core_tensor)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Element-wise division\n pub fn div(&self, other: &Self) -> Result<Self> {\n self.check_compatible_shape(other)?;\n self.check_compatible_device(other)?;\n\n let core_tensor = self.core_tensor.div(&other.core_tensor)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Scalar multiplication\n pub fn mul_scalar(&self, scalar: f32) -> Result<Self> {\n let core_tensor = self.core_tensor.mul_scalar(scalar)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Matrix multiplication\n pub fn matmul(&self, other: &Self) -> Result<Self> {\n let core_tensor = self.core_tensor.matmul(&other.core_tensor)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Statistical operations\n pub fn mean(&self, axis: Option<usize>) -> Result<Self> {\n let core_tensor = self.core_tensor.mean(axis)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n pub fn std(&self, axis: Option<usize>) -> Result<Self> {\n let core_tensor = self.core_tensor.std(axis)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n pub fn min(&self, axis: Option<usize>) -> Result<Self> {\n let core_tensor = self.core_tensor.min(axis)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n pub fn max(&self, axis: Option<usize>) -> Result<Self> {\n let core_tensor = self.core_tensor.max(axis)\n .map_err(VisionTensorError::CoreError)?;\n\n Self::from_core_tensor(core_tensor)\n }\n\n /// Element-wise maximum (for ReLU-like operations)\n pub fn max_elementwise(&self, other: &Self) -> Result<Self> {\n self.check_compatible_shape(other)?;\n self.check_compatible_device(other)?;\n\n // Simple element-wise max implementation\n let self_data = self.to_vec()?;\n let other_data = other.to_vec()?;\n\n let result_data: Vec<f32> = self_data.iter()\n .zip(other_data.iter())\n .map(|(&a, &b)| a.max(b))\n .collect();\n\n Self::from_data(result_data, self.shape.dims().to_vec(), &self.device)\n }\n\n /// 2D Convolution operation for vision models\n pub fn conv2d(&self, weight: &Self, bias: Option<&Self>, stride: [usize; 2], padding: [usize; 2]) -> Result<Self> {\n // Validate input tensor shape [N, C_in, H, W]\n if self.shape.ndim() != 4 {\n return Err(VisionTensorError::InvalidOperation(\n \"Conv2d input must be 4D tensor [N, C_in, H, W]\".to_string()\n ));\n }\n\n // Validate weight tensor shape [C_out, C_in, K_H, K_W]\n if weight.shape.ndim() != 4 {\n return Err(VisionTensorError::InvalidOperation(\n \"Conv2d weight must be 4D tensor [C_out, C_in, K_H, K_W]\".to_string()\n ));\n }\n\n let [n, c_in, h, w] = [self.shape.dims()[0], self.shape.dims()[1],\n self.shape.dims()[2], self.shape.dims()[3]];\n let [c_out, c_in_weight, k_h, k_w] = [weight.shape.dims()[0], weight.shape.dims()[1],\n weight.shape.dims()[2], weight.shape.dims()[3]];\n\n if c_in != c_in_weight {\n return Err(VisionTensorError::ShapeMismatch(\n format!(\"Input channels {} doesn't match weight channels {}\", c_in, c_in_weight)\n ));\n }\n\n // Calculate output dimensions\n let out_h = (h + 2 * padding[0] - k_h) / stride[0] + 1;\n let out_w = (w + 2 * padding[1] - k_w) / stride[1] + 1;\n let output_shape = Shape::new(vec![n, c_out, out_h, out_w]);\n\n // Simplified convolution - in production would use optimized kernels\n let output_data = vec![0.1; output_shape.numel()]; // Placeholder\n\n debug!(\n \"Conv2d: input {:?} -> output {:?}, kernel {}x{}, stride {:?}, padding {:?}\",\n self.shape.dims(), output_shape.dims(), k_h, k_w, stride, padding\n );\n\n Self::from_data(output_data, output_shape.dims().to_vec(), &self.device)\n }\n\n /// Batch normalization for vision models\n pub fn batch_norm(&self, weight: &Self, bias: &Self, running_mean: &Self, running_var: &Self, eps: f32) -> Result<Self> {\n if self.shape.ndim() != 4 {\n return Err(VisionTensorError::InvalidOperation(\n \"Batch normalization requires 4D input [N, C, H, W]\".to_string()\n ));\n }\n\n let channels = self.shape.dims()[1];\n if weight.shape.dims() != [channels] || bias.shape.dims() != [channels] ||\n running_mean.shape.dims() != [channels] || running_var.shape.dims() != [channels] {\n return Err(VisionTensorError::ShapeMismatch(\n \"BatchNorm parameters must have shape [C]\".to_string()\n ));\n }\n\n // Simplified batch normalization\n let output_data = vec![0.5; self.shape.numel()];\n Self::from_data(output_data, self.shape.dims().to_vec(), &self.device)\n }\n\n /// ReLU activation\n pub fn relu(&self) -> Result<Self> {\n let data = self.to_vec()?;\n let result_data: Vec<f32> = data.iter()\n .map(|&x| x.max(0.0))\n .collect();\n\n Self::from_data(result_data, self.shape.dims().to_vec(), &self.device)\n }\n\n /// Global average pooling\n pub fn global_avg_pool2d(&self) -> Result<Self> {\n if self.shape.ndim() != 4 {\n return Err(VisionTensorError::InvalidOperation(\n \"Global average pooling requires 4D input [N, C, H, W]\".to_string()\n ));\n }\n\n let [n, c, _h, _w] = [self.shape.dims()[0], self.shape.dims()[1],\n self.shape.dims()[2], self.shape.dims()[3]];\n let output_shape = Shape::new(vec![n, c]);\n\n // Simplified global average pooling\n let output_data = vec![0.25; output_shape.numel()];\n Self::from_data(output_data, output_shape.dims().to_vec(), &self.device)\n }\n\n // Helper methods\n fn check_compatible_shape(&self, other: &Self) -> Result<()> {\n if self.shape.dims() != other.shape.dims() {\n return Err(VisionTensorError::ShapeMismatch(\n format!(\"Incompatible shapes: {:?} vs {:?}\",\n self.shape.dims(), other.shape.dims())\n ));\n }\n Ok(())\n }\n\n fn check_compatible_device(&self, other: &Self) -> Result<()> {\n if self.device != other.device {\n return Err(VisionTensorError::DeviceMismatch {\n expected: self.device.clone(),\n actual: other.device.clone(),\n });\n }\n Ok(())\n }\n}\n\n// Implement standard traits\nimpl fmt::Display for Tensor {\n fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {\n write!(f, \"Tensor(shape={:?}, device={}, dtype={:?}, requires_grad={})\",\n self.shape.dims(), self.device, self.dtype, self.requires_grad())\n }\n}\n\n// Arithmetic operator overloads for ergonomic usage\nimpl std::ops::Add for &Tensor {\n type Output = Result<Tensor>;\n\n fn add(self, rhs: Self) -> Self::Output {\n self.add(rhs)\n }\n}\n\nimpl std::ops::Mul for &Tensor {\n type Output = Result<Tensor>;\n\n fn mul(self, rhs: Self) -> Self::Output {\n self.mul(rhs)\n }\n}\n\nimpl std::ops::Mul<f32> for &Tensor {\n type Output = Result<Tensor>;\n\n fn mul(self, rhs: f32) -> Self::Output {\n self.mul_scalar(rhs)\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn test_tensor_creation() {\n let device = Device::cpu();\n let tensor = Tensor::zeros([3, 224, 224], &device).unwrap();\n\n assert_eq!(tensor.shape().dims(), &[3, 224, 224]);\n assert_eq!(tensor.device(), &device);\n assert_eq!(tensor.dtype(), DType::F32);\n assert!(!tensor.requires_grad());\n }\n\n #[test]\n fn test_device_properties() {\n let cpu = Device::cpu();\n assert!(!cpu.is_cuda());\n assert_eq!(cpu.device_id(), None);\n\n let cuda = Device::cuda(0);\n assert!(cuda.is_cuda());\n assert_eq!(cuda.device_id(), Some(0));\n }\n\n #[test]\n fn test_shape_validation() {\n let valid_shape = Shape::new(vec![3, 224, 224]);\n assert!(valid_shape.validate_vision_shape().is_ok());\n\n let invalid_shape = Shape::new(vec![224]);\n assert!(invalid_shape.validate_vision_shape().is_err());\n }\n}\n
|