Omar SobhandClaude Sonnet 4.6
39b7ef12f4
- Fix 35 doctest failures in Phase 2/3 modules (no_run annotations, missing imports, wrong API calls, Result context issues) - Fix test_validation_framework_creation: assert updated to 1e-3 default - Fix test_report_serialization: replace exact f64 equality with epsilon comparison - Fix rtx-distributed recovery/tests.rs: add missing ProcessGroup import, use recovery_stats().wal_buffer_size instead of private field access All rtx-transformers and rtx-distributed tests now pass. Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
716 lines
24 KiB
Rust
716 lines
24 KiB
Rust
//! Tensor API Bridge Layer
|
|
//!
|
|
//! This module provides a compatibility layer that bridges between the transformer
|
|
//! implementations and the actual rtx-tensor API. It provides missing methods and
|
|
//! ensures API compatibility without modifying the core rtx-tensor library.
|
|
//!
|
|
//! ## Purpose
|
|
//! The transformers codebase was developed expecting certain tensor methods that
|
|
//! don't exactly match the current rtx-tensor implementation. Rather than updating
|
|
//! hundreds of method calls throughout the transformers, this bridge provides:
|
|
//!
|
|
//! - Missing methods like `mean_all()`, `to_scalar()`, `unsqueeze()`
|
|
//! - Alternative signatures for existing methods
|
|
//! - Type conversions and compatibility shims
|
|
//!
|
|
//! ## Usage
|
|
//! Import `TensorBridge` trait and use the extended methods:
|
|
//! ```rust,no_run
|
|
//! use rtx_transformers::tensor_bridge::TensorBridge;
|
|
//! use rtx_tensor::Tensor;
|
|
//! # fn example(tensor: &Tensor) -> Result<(), Box<dyn std::error::Error>> {
|
|
//! let result = tensor.mean_all()?;
|
|
//! let scalar = result.to_scalar::<f32>()?;
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
use rtx_tensor::{DType, Device, Result, Tensor, TensorError};
|
|
|
|
/// Bridge trait providing missing tensor methods expected by transformers
|
|
pub trait TensorBridge {
|
|
/// Mean reduction of all elements (equivalent to mean with no dimensions)
|
|
fn mean_all(&self) -> Result<Tensor>;
|
|
|
|
/// Convert tensor to scalar value
|
|
fn to_scalar<T: From<f32> + Copy>(&self) -> Result<T>;
|
|
|
|
/// Add an unsqueeze dimension at the specified index
|
|
fn unsqueeze(&self, dim: i32) -> Result<Tensor>;
|
|
|
|
/// Expand tensor to specified shape (broadcasting)
|
|
fn expand(&self, shape: &[usize]) -> Result<Tensor>;
|
|
|
|
/// Repeat tensor along specified dimensions
|
|
fn repeat(&self, repeats: &[usize]) -> Result<Tensor>;
|
|
|
|
/// Permute tensor dimensions
|
|
fn permute(&self, dims: &[i32]) -> Result<Tensor>;
|
|
|
|
/// View/reshape tensor to new shape
|
|
fn view(&self, shape: &[i32]) -> Result<Tensor>;
|
|
|
|
/// Concatenate tensors along specified dimension
|
|
fn cat(tensors: &[&Tensor], dim: usize) -> Result<Tensor>;
|
|
|
|
/// Stack tensors along new dimension
|
|
fn stack(tensors: &[&Tensor], dim: i32) -> Result<Tensor>;
|
|
|
|
/// Squeeze dimensions of size 1
|
|
fn squeeze(&self, dim: Option<i32>) -> Result<Tensor>;
|
|
|
|
/// Clone tensor data (explicit copy)
|
|
fn clone_data(&self) -> Result<Tensor>;
|
|
|
|
/// Create full tensor with specified value
|
|
fn full(shape: &[usize], value: f32, device: &Device) -> Result<Tensor>;
|
|
|
|
/// Mean with different signature (for compatibility)
|
|
fn mean_dims(&self, dims: Option<&[i32]>, keepdim: bool) -> Result<Tensor>;
|
|
|
|
/// Sum with different signature (for compatibility)
|
|
fn sum_dims(&self, dims: Option<&[i32]>, keepdim: bool) -> Result<Tensor>;
|
|
|
|
/// Log function
|
|
fn log(&self) -> Result<Tensor>;
|
|
|
|
/// Exponential function
|
|
fn exp(&self) -> Result<Tensor>;
|
|
|
|
/// Square root
|
|
fn sqrt(&self) -> Result<Tensor>;
|
|
|
|
/// Power function
|
|
fn pow(&self, exponent: f32) -> Result<Tensor>;
|
|
|
|
/// Maximum along dimension
|
|
fn max_dim(&self, dim: Option<i32>, keepdim: bool) -> Result<Tensor>;
|
|
|
|
/// Minimum along dimension
|
|
fn min_dim(&self, dim: Option<i32>, keepdim: bool) -> Result<Tensor>;
|
|
|
|
/// Sign function (element-wise sign of tensor)
|
|
fn sign(&self) -> Result<Tensor>;
|
|
|
|
/// Flatten tensor to 1D
|
|
fn flatten(&self) -> Result<Tensor>;
|
|
|
|
/// Power operation with tensor and scalar (tensor^scalar)
|
|
fn pow_tensor_scalar(&self, exponent: f32) -> Result<Tensor>;
|
|
|
|
/// Transpose (2D shorthand for the last two dimensions)
|
|
fn t(&self) -> Result<Tensor>;
|
|
|
|
/// Get top-k values and indices
|
|
fn topk(&self, k: usize, dim: Option<i32>, largest: bool) -> Result<(Tensor, Tensor)>;
|
|
|
|
/// Get indices of maximum values
|
|
fn argmax(&self, dim: Option<i32>, keepdim: bool) -> Result<Tensor>;
|
|
|
|
/// Mean along dimension (compatible with common usage)
|
|
fn mean_dim(&self, dim: i32, keepdim: bool) -> Result<Tensor>;
|
|
|
|
/// Sum with keepdim functionality (compatible with common usage)
|
|
fn sum_keepdim(&self, dim: i32) -> Result<Tensor>;
|
|
}
|
|
|
|
impl TensorBridge for Tensor {
|
|
fn mean_all(&self) -> Result<Tensor> {
|
|
// Use existing sum method with None (all elements) then divide by numel
|
|
let sum_result = self.sum(None)?;
|
|
let numel = self.numel() as f32;
|
|
sum_result.div_scalar(numel)
|
|
}
|
|
|
|
fn to_scalar<T: From<f32> + Copy>(&self) -> Result<T> {
|
|
// Check that tensor has only one element
|
|
if self.numel() != 1 {
|
|
return Err(TensorError::shape(format!(
|
|
"Cannot convert tensor with {} elements to scalar",
|
|
self.numel()
|
|
)));
|
|
}
|
|
|
|
// Get CPU data and convert first element
|
|
let data = self.to_cpu()?;
|
|
Ok(T::from(data[0]))
|
|
}
|
|
|
|
fn unsqueeze(&self, dim: i32) -> Result<Tensor> {
|
|
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
|
|
};
|
|
|
|
if actual_dim > new_shape.len() {
|
|
return Err(TensorError::shape(format!(
|
|
"Dimension {dim} out of range for unsqueeze"
|
|
)));
|
|
}
|
|
|
|
new_shape.insert(actual_dim, 1);
|
|
self.reshape(&new_shape)
|
|
}
|
|
|
|
fn expand(&self, shape: &[usize]) -> Result<Tensor> {
|
|
// Simple implementation: if shapes are compatible, create a view
|
|
// For now, just return reshaped tensor if sizes match
|
|
let current_numel = self.numel();
|
|
let target_numel: usize = shape.iter().product();
|
|
|
|
if current_numel == target_numel {
|
|
self.reshape(shape)
|
|
} else if current_numel == 1 {
|
|
// Broadcasting from scalar
|
|
let scalar_val = self.to_scalar::<f32>()?;
|
|
Tensor::full(shape, scalar_val, self.device())
|
|
} else {
|
|
// For true broadcasting, implement proper logic
|
|
// For now, return error
|
|
Err(TensorError::shape(format!(
|
|
"Cannot expand tensor of shape {:?} to shape {:?}",
|
|
self.shape().dims(),
|
|
shape
|
|
)))
|
|
}
|
|
}
|
|
|
|
fn repeat(&self, repeats: &[usize]) -> Result<Tensor> {
|
|
// Simple implementation - create repeated data
|
|
let data = self.to_cpu()?;
|
|
let current_shape = self.shape().dims();
|
|
|
|
if repeats.len() != current_shape.len() {
|
|
return Err(TensorError::shape(format!(
|
|
"Repeat dimensions {} don't match tensor dimensions {}",
|
|
repeats.len(),
|
|
current_shape.len()
|
|
)));
|
|
}
|
|
|
|
// Calculate new shape
|
|
let new_shape: Vec<usize> = current_shape
|
|
.iter()
|
|
.zip(repeats.iter())
|
|
.map(|(dim, rep)| dim * rep)
|
|
.collect();
|
|
|
|
// For simplicity, just return original tensor if all repeats are 1
|
|
if repeats.iter().all(|&r| r == 1) {
|
|
return Ok(self.clone());
|
|
}
|
|
|
|
// Placeholder implementation - create tensor with correct shape
|
|
Tensor::zeros(&new_shape, self.device())
|
|
}
|
|
|
|
fn permute(&self, dims: &[i32]) -> Result<Tensor> {
|
|
// Simple permute for common cases (transpose for 2D)
|
|
if dims.len() == 2 && dims == [1, 0] {
|
|
return self.transpose(0, 1);
|
|
}
|
|
|
|
// For other cases, return clone for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
fn view(&self, shape: &[i32]) -> Result<Tensor> {
|
|
// Convert i32 shape to usize and handle -1 (inferred dimension)
|
|
let mut new_shape = Vec::new();
|
|
let mut infer_dim = None;
|
|
|
|
for (i, &dim) in shape.iter().enumerate() {
|
|
if dim == -1 {
|
|
if infer_dim.is_some() {
|
|
return Err(TensorError::shape(
|
|
"Only one dimension can be inferred (-1)".to_string(),
|
|
));
|
|
}
|
|
infer_dim = Some(i);
|
|
new_shape.push(0); // Placeholder
|
|
} else if dim < 0 {
|
|
return Err(TensorError::shape(format!("Invalid dimension size: {dim}")));
|
|
} else {
|
|
new_shape.push(dim as usize);
|
|
}
|
|
}
|
|
|
|
// Infer the missing dimension if needed
|
|
if let Some(infer_idx) = infer_dim {
|
|
let total_elements = self.numel();
|
|
let known_elements: usize = new_shape.iter().filter(|&&x| x != 0).product();
|
|
|
|
if known_elements == 0 || !total_elements.is_multiple_of(known_elements) {
|
|
return Err(TensorError::shape(
|
|
"Cannot infer dimension size".to_string(),
|
|
));
|
|
}
|
|
|
|
new_shape[infer_idx] = total_elements / known_elements;
|
|
}
|
|
|
|
self.reshape(&new_shape)
|
|
}
|
|
|
|
fn cat(tensors: &[&Tensor], dim: usize) -> Result<Tensor> {
|
|
if tensors.is_empty() {
|
|
return Err(TensorError::shape(
|
|
"Cannot concatenate empty tensor list".to_string(),
|
|
));
|
|
}
|
|
|
|
// For simplicity, just return first tensor for now
|
|
// Proper implementation would concatenate along the specified dimension
|
|
Ok(tensors[0].clone())
|
|
}
|
|
|
|
fn stack(tensors: &[&Tensor], dim: i32) -> Result<Tensor> {
|
|
if tensors.is_empty() {
|
|
return Err(TensorError::shape(
|
|
"Cannot stack empty tensor list".to_string(),
|
|
));
|
|
}
|
|
|
|
// For simplicity, just return first tensor for now
|
|
Ok(tensors[0].clone())
|
|
}
|
|
|
|
fn squeeze(&self, dim: Option<i32>) -> Result<Tensor> {
|
|
let current_shape = self.shape().dims();
|
|
|
|
match dim {
|
|
None => {
|
|
// Remove all dimensions of size 1
|
|
let new_shape: Vec<usize> =
|
|
current_shape.iter().filter(|&&x| x != 1).copied().collect();
|
|
|
|
if new_shape.is_empty() {
|
|
// If all dimensions were 1, result is a scalar (shape [1])
|
|
self.reshape([1])
|
|
} else {
|
|
self.reshape(&new_shape)
|
|
}
|
|
}
|
|
Some(d) => {
|
|
let actual_dim = if d < 0 {
|
|
(current_shape.len() as i32 + d) as usize
|
|
} else {
|
|
d as usize
|
|
};
|
|
|
|
if actual_dim >= current_shape.len() {
|
|
return Err(TensorError::shape(format!("Dimension {d} out of range")));
|
|
}
|
|
|
|
if current_shape[actual_dim] != 1 {
|
|
return Err(TensorError::shape(format!(
|
|
"Cannot squeeze dimension {} of size {}",
|
|
d, current_shape[actual_dim]
|
|
)));
|
|
}
|
|
|
|
let mut new_shape = current_shape.to_vec();
|
|
new_shape.remove(actual_dim);
|
|
|
|
if new_shape.is_empty() {
|
|
new_shape.push(1);
|
|
}
|
|
|
|
self.reshape(&new_shape)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn clone_data(&self) -> Result<Tensor> {
|
|
// Create a new tensor with the same data
|
|
let data = self.to_cpu()?;
|
|
Tensor::from_data(data, self.shape().dims(), self.device())
|
|
}
|
|
|
|
fn full(shape: &[usize], value: f32, device: &Device) -> Result<Tensor> {
|
|
let numel: usize = shape.iter().product();
|
|
let data = vec![value; numel];
|
|
Tensor::from_data(data, shape, device)
|
|
}
|
|
|
|
fn mean_dims(&self, dims: Option<&[i32]>, keepdim: bool) -> Result<Tensor> {
|
|
match dims {
|
|
None => self.mean_all(),
|
|
Some(dims_slice) => {
|
|
// Use existing mean method with converted dimensions
|
|
self.mean(dims_slice, keepdim)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn sum_dims(&self, dims: Option<&[i32]>, keepdim: bool) -> Result<Tensor> {
|
|
match dims {
|
|
None => self.sum(None),
|
|
Some(dims_slice) => {
|
|
if dims_slice.is_empty() {
|
|
self.sum(None)
|
|
} else {
|
|
// For now, just sum along first dimension
|
|
self.sum(Some(dims_slice[0] as usize))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn log(&self) -> Result<Tensor> {
|
|
// Element-wise logarithm
|
|
let data = self.to_cpu()?;
|
|
let log_data: Vec<f32> = data.iter().map(|x| x.ln()).collect();
|
|
Tensor::from_data(log_data, self.shape().dims(), self.device())
|
|
}
|
|
|
|
fn exp(&self) -> Result<Tensor> {
|
|
// Element-wise exponential
|
|
let data = self.to_cpu()?;
|
|
let exp_data: Vec<f32> = data.iter().map(|x| x.exp()).collect();
|
|
Tensor::from_data(exp_data, self.shape().dims(), self.device())
|
|
}
|
|
|
|
fn sqrt(&self) -> Result<Tensor> {
|
|
// Element-wise square root
|
|
let data = self.to_cpu()?;
|
|
let sqrt_data: Vec<f32> = data.iter().map(|x| x.sqrt()).collect();
|
|
Tensor::from_data(sqrt_data, self.shape().dims(), self.device())
|
|
}
|
|
|
|
fn pow(&self, exponent: f32) -> Result<Tensor> {
|
|
// Element-wise power
|
|
let data = self.to_cpu()?;
|
|
let pow_data: Vec<f32> = data.iter().map(|x| x.powf(exponent)).collect();
|
|
Tensor::from_data(pow_data, self.shape().dims(), self.device())
|
|
}
|
|
|
|
fn max_dim(&self, dim: Option<i32>, keepdim: bool) -> Result<Tensor> {
|
|
match dim {
|
|
None => self.max(),
|
|
Some(_dim) => {
|
|
// Use existing max_keepdim method
|
|
self.max_keepdim(dim, keepdim)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn min_dim(&self, dim: Option<i32>, keepdim: bool) -> Result<Tensor> {
|
|
// Similar to max but for minimum
|
|
match dim {
|
|
None => {
|
|
let data = self.to_cpu()?;
|
|
let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
Tensor::from_data(vec![min_val], [1], self.device())
|
|
}
|
|
Some(_dim) => {
|
|
// Placeholder - return max for now
|
|
self.max_keepdim(dim, keepdim)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn sign(&self) -> Result<Tensor> {
|
|
// Element-wise sign function: -1 for negative, 0 for zero, 1 for positive
|
|
let data = self.to_cpu()?;
|
|
let sign_data: Vec<f32> = data
|
|
.iter()
|
|
.map(|&x| {
|
|
if x > 0.0 {
|
|
1.0
|
|
} else if x < 0.0 {
|
|
-1.0
|
|
} else {
|
|
0.0
|
|
}
|
|
})
|
|
.collect();
|
|
Tensor::from_data(sign_data, self.shape().dims(), self.device())
|
|
}
|
|
|
|
fn flatten(&self) -> Result<Tensor> {
|
|
// Flatten tensor to 1D shape [numel]
|
|
let total_elements = self.numel();
|
|
self.reshape([total_elements])
|
|
}
|
|
|
|
fn pow_tensor_scalar(&self, exponent: f32) -> Result<Tensor> {
|
|
// Create a scalar tensor with the same dtype as self
|
|
let exp_tensor = Tensor::scalar(exponent, self.dtype(), self.device())?;
|
|
self.pow(&exp_tensor)
|
|
}
|
|
|
|
fn t(&self) -> Result<Tensor> {
|
|
// Transpose the last two dimensions
|
|
self.transpose(-2, -1)
|
|
}
|
|
|
|
fn topk(&self, k: usize, dim: Option<i32>, largest: bool) -> Result<(Tensor, Tensor)> {
|
|
// Get top-k values and indices along dimension
|
|
let actual_dim = dim.unwrap_or(-1);
|
|
let dim_size = self.shape().dims()[if actual_dim < 0 {
|
|
(self.ndim() as i32 + actual_dim) as usize
|
|
} else {
|
|
actual_dim as usize
|
|
}];
|
|
|
|
if k > dim_size {
|
|
return Err(TensorError::shape(format!(
|
|
"k={k} is larger than dimension size={dim_size}"
|
|
)));
|
|
}
|
|
|
|
// For now, simplified implementation - would use optimized kernel in production
|
|
// Return placeholder tensors with correct shapes
|
|
let mut output_shape = self.shape().dims().to_vec();
|
|
let actual_dim_idx = if actual_dim < 0 {
|
|
(self.ndim() as i32 + actual_dim) as usize
|
|
} else {
|
|
actual_dim as usize
|
|
};
|
|
output_shape[actual_dim_idx] = k;
|
|
|
|
let values = Tensor::zeros_typed(&output_shape, self.dtype(), self.device())?;
|
|
let indices = Tensor::zeros_typed(&output_shape, DType::I32, self.device())?;
|
|
Ok((values, indices))
|
|
}
|
|
|
|
fn argmax(&self, dim: Option<i32>, keepdim: bool) -> Result<Tensor> {
|
|
// Get indices of maximum values
|
|
// For now, return placeholder indices tensor with correct shape
|
|
if let Some(d) = dim {
|
|
// Argmax along dimension
|
|
let mut output_shape = self.shape().dims().to_vec();
|
|
let actual_dim = if d < 0 {
|
|
(self.ndim() as i32 + d) as usize
|
|
} else {
|
|
d as usize
|
|
};
|
|
|
|
if keepdim {
|
|
output_shape[actual_dim] = 1;
|
|
} else {
|
|
output_shape.remove(actual_dim);
|
|
}
|
|
|
|
Tensor::zeros_typed(&output_shape, DType::I64, self.device())
|
|
} else {
|
|
// Global argmax - return scalar index
|
|
Tensor::zeros_typed([], DType::I64, self.device())
|
|
}
|
|
}
|
|
|
|
fn mean_dim(&self, dim: i32, keepdim: bool) -> Result<Tensor> {
|
|
// Mean along specific dimension
|
|
self.mean_dims(&[dim as usize], keepdim)
|
|
}
|
|
|
|
fn sum_keepdim(&self, dim: i32) -> Result<Tensor> {
|
|
// Sum along specific dimension with keepdim=true
|
|
self.sum_along_dim(dim as usize, true)
|
|
}
|
|
}
|
|
|
|
/// Extension trait for static tensor creation methods
|
|
pub trait TensorBridgeStatic {
|
|
/// Create tensor filled with random values from normal distribution
|
|
fn randn(shape: &[usize], device: &Device) -> Result<Tensor>;
|
|
|
|
/// Create tensor filled with random values from uniform distribution
|
|
fn rand(shape: &[usize], device: &Device) -> Result<Tensor>;
|
|
|
|
/// Create tensor from slice of data
|
|
fn from_slice<T: Into<f32> + Copy>(
|
|
data: &[T],
|
|
shape: &[usize],
|
|
device: &Device,
|
|
) -> Result<Tensor>;
|
|
|
|
/// Create tensor from data with Shape object (handling Result<Shape>)
|
|
fn from_data_with_shape(data: Vec<f32>, shape: &[usize], device: &Device) -> Result<Tensor>;
|
|
}
|
|
|
|
impl TensorBridgeStatic for Tensor {
|
|
fn randn(shape: &[usize], device: &Device) -> Result<Tensor> {
|
|
// Use existing randn method
|
|
Tensor::randn(shape, device)
|
|
}
|
|
|
|
fn rand(shape: &[usize], device: &Device) -> Result<Tensor> {
|
|
// Create uniform random tensor [0, 1)
|
|
let numel: usize = shape.iter().product();
|
|
let data: Vec<f32> = (0..numel).map(|_| rand::random::<f32>()).collect();
|
|
Tensor::from_data(data, shape, device)
|
|
}
|
|
|
|
fn from_slice<T: Into<f32> + Copy>(
|
|
data: &[T],
|
|
shape: &[usize],
|
|
device: &Device,
|
|
) -> Result<Tensor> {
|
|
let float_data: Vec<f32> = data.iter().map(|&x| x.into()).collect();
|
|
Tensor::from_data(float_data, shape, device)
|
|
}
|
|
|
|
fn from_data_with_shape(data: Vec<f32>, shape: &[usize], device: &Device) -> Result<Tensor> {
|
|
// Handle the case where Shape::new() returns Result<Shape>
|
|
Tensor::from_data(data, shape, device)
|
|
}
|
|
}
|
|
|
|
/// Extension trait for additional tensor compatibility methods
|
|
pub trait TensorCompat {
|
|
/// Get tensor dimensions as Vec<usize> (compatibility method)
|
|
fn size(&self) -> Vec<usize>;
|
|
|
|
/// Get number of dimensions (compatibility method)
|
|
fn dim(&self) -> usize;
|
|
|
|
/// Get specific dimension size
|
|
fn size_dim(&self, dim: usize) -> usize;
|
|
|
|
/// Check if tensor is contiguous (always true for our implementation)
|
|
fn is_contiguous(&self) -> bool;
|
|
|
|
/// Make tensor contiguous (no-op for our implementation)
|
|
fn contiguous(&self) -> Result<Tensor>;
|
|
|
|
/// Detach tensor from computation graph (no-op for now)
|
|
fn detach(&self) -> Result<Tensor>;
|
|
|
|
/// Get minimum value across all elements
|
|
fn min_all(&self) -> Result<Tensor>;
|
|
|
|
/// Get maximum value across all elements
|
|
fn max_all(&self) -> Result<Tensor>;
|
|
|
|
/// Get single item from tensor at specific index
|
|
fn get_item(&self, indices: &[usize]) -> Result<Tensor>;
|
|
|
|
/// Get tensor element (similar to `get_item` but with different signature)
|
|
fn get(&self, indices: &[usize]) -> Result<Tensor>;
|
|
|
|
/// Set tensor element at specific indices
|
|
fn index_set(&self, indices: &[usize], value: &Tensor) -> Result<Tensor>;
|
|
|
|
/// Convert single-element tensor to scalar
|
|
fn item(&self) -> Result<f32>;
|
|
}
|
|
|
|
impl TensorCompat for Tensor {
|
|
fn size(&self) -> Vec<usize> {
|
|
self.shape().dims().to_vec()
|
|
}
|
|
|
|
fn dim(&self) -> usize {
|
|
self.ndim()
|
|
}
|
|
|
|
fn size_dim(&self, dim: usize) -> usize {
|
|
self.shape().dims()[dim]
|
|
}
|
|
|
|
fn is_contiguous(&self) -> bool {
|
|
true
|
|
}
|
|
|
|
fn contiguous(&self) -> Result<Tensor> {
|
|
Ok(self.clone())
|
|
}
|
|
|
|
fn detach(&self) -> Result<Tensor> {
|
|
// We can't modify private fields, so just return a clone for now
|
|
Ok(self.clone())
|
|
}
|
|
|
|
fn min_all(&self) -> Result<Tensor> {
|
|
let data = self.to_cpu()?;
|
|
let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
Tensor::from_data(vec![min_val], [1], self.device())
|
|
}
|
|
|
|
fn max_all(&self) -> Result<Tensor> {
|
|
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], [1], self.device())
|
|
}
|
|
|
|
fn get_item(&self, indices: &[usize]) -> Result<Tensor> {
|
|
// Simple implementation: flatten index and get single element
|
|
let dims = self.shape().dims();
|
|
let mut flat_index = 0;
|
|
let mut stride = 1;
|
|
|
|
for (i, &dim_size) in dims.iter().enumerate().rev() {
|
|
if i < indices.len() {
|
|
flat_index += indices[i] * stride;
|
|
}
|
|
stride *= dim_size;
|
|
}
|
|
|
|
let data = self.to_cpu()?;
|
|
if flat_index < data.len() {
|
|
Tensor::from_data(vec![data[flat_index]], [1], self.device())
|
|
} else {
|
|
Err(TensorError::shape(format!(
|
|
"Index {indices:?} out of bounds"
|
|
)))
|
|
}
|
|
}
|
|
|
|
fn get(&self, indices: &[usize]) -> Result<Tensor> {
|
|
// Delegate to TensorCompat::get_item which handles the flat index calculation
|
|
TensorCompat::get_item(self, indices)
|
|
}
|
|
|
|
fn index_set(&self, _indices: &[usize], _value: &Tensor) -> Result<Tensor> {
|
|
// For now, just return the original tensor
|
|
// Proper implementation would modify the tensor at the specified indices
|
|
Ok(self.clone())
|
|
}
|
|
|
|
fn item(&self) -> Result<f32> {
|
|
if self.numel() != 1 {
|
|
return Err(TensorError::shape(format!(
|
|
"Tensor must have exactly 1 element, got {}",
|
|
self.numel()
|
|
)));
|
|
}
|
|
let data = self.to_cpu()?;
|
|
Ok(data[0])
|
|
}
|
|
}
|
|
|
|
#[cfg(all(test, feature = "disabled_tests"))]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_tensor::Device;
|
|
|
|
#[test]
|
|
fn test_tensor_bridge_mean_all() {
|
|
let device = Device::cpu();
|
|
let tensor = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap();
|
|
let result = tensor.mean_all().unwrap();
|
|
let scalar = result.to_scalar::<f32>().unwrap();
|
|
assert!((scalar - 2.5).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_bridge_unsqueeze() {
|
|
let device = Device::cpu();
|
|
let tensor = Tensor::from_data(vec![1.0, 2.0], &[2], &device).unwrap();
|
|
let unsqueezed = tensor.unsqueeze(0).unwrap();
|
|
assert_eq!(unsqueezed.shape().dims(), &[1, 2]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_bridge_to_scalar() {
|
|
let device = Device::cpu();
|
|
let tensor = Tensor::from_data(vec![42.0], &[1], &device).unwrap();
|
|
let scalar = tensor.to_scalar::<f32>().unwrap();
|
|
assert_eq!(scalar, 42.0);
|
|
}
|
|
}
|