129 lines
4.0 KiB
Rust
129 lines
4.0 KiB
Rust
//! Utility functions for tensor operations that aren't available in rtx-tensor
|
|
//! Following strict TDD approach - implementing missing functionality
|
|
//!
|
|
//! This module includes tensor indexing utilities needed for cross-validation splits.
|
|
|
|
use crate::{AutoMLError, AutoMLResult};
|
|
use rtx_tensor::{Result, Tensor};
|
|
|
|
/// Find minimum value in a tensor
|
|
/// Since rtx-tensor doesn't have min(), we implement it using max with negation
|
|
pub fn find_min(tensor: &Tensor) -> Result<f32> {
|
|
// Negate the tensor, find max, then negate result
|
|
let negated = tensor.mul_scalar(-1.0)?;
|
|
let max_of_negated = negated.max()?.to_scalar::<f32>()?;
|
|
Ok(-max_of_negated)
|
|
}
|
|
|
|
/// Find minimum value in a tensor and return as Tensor
|
|
pub fn tensor_min(tensor: &Tensor) -> Result<Tensor> {
|
|
let min_val = find_min(tensor)?;
|
|
Tensor::scalar(min_val, tensor.dtype(), tensor.device())
|
|
}
|
|
|
|
/// Index a tensor by a list of sample indices (for CV train/test splits)
|
|
///
|
|
/// This function selects rows from a tensor based on the provided indices,
|
|
/// useful for creating train/validation splits from cross-validation.
|
|
///
|
|
/// # Arguments
|
|
/// * `tensor` - The input tensor with shape [n_samples, ...]
|
|
/// * `indices` - The indices of samples to select
|
|
///
|
|
/// # Returns
|
|
/// A new tensor containing only the selected samples
|
|
///
|
|
/// # Example
|
|
/// ```ignore
|
|
/// let x = Tensor::randn(&[100, 10], &device)?;
|
|
/// let train_idx = vec![0, 2, 4, 6, 8];
|
|
/// let x_train = index_tensor(&x, &train_idx)?;
|
|
/// // x_train has shape [5, 10]
|
|
/// ```
|
|
pub fn index_tensor(tensor: &Tensor, indices: &[usize]) -> AutoMLResult<Tensor> {
|
|
if indices.is_empty() {
|
|
return Err(AutoMLError::ValidationError(
|
|
"Cannot index tensor with empty indices".to_string(),
|
|
));
|
|
}
|
|
|
|
let shape = tensor.shape();
|
|
let n_samples = shape.dims()[0];
|
|
|
|
// Validate indices
|
|
for &idx in indices {
|
|
if idx >= n_samples {
|
|
return Err(AutoMLError::ValidationError(format!(
|
|
"Index {idx} out of bounds for tensor with {n_samples} samples"
|
|
)));
|
|
}
|
|
}
|
|
|
|
// Use index_select on dimension 0 (samples dimension)
|
|
// rtx-tensor's index_select takes &[usize] directly
|
|
tensor
|
|
.index_select(0, indices)
|
|
.map_err(|e| AutoMLError::ModelError(e.to_string()))
|
|
}
|
|
|
|
/// Split tensor data into train and validation sets based on indices
|
|
///
|
|
/// # Arguments
|
|
/// * `x` - Features tensor with shape [n_samples, n_features]
|
|
/// * `y` - Labels tensor with shape [n_samples] or [n_samples, n_outputs]
|
|
/// * `train_idx` - Indices for training samples
|
|
/// * `val_idx` - Indices for validation samples
|
|
///
|
|
/// # Returns
|
|
/// Tuple of (x_train, y_train, x_val, y_val)
|
|
pub fn split_train_val(
|
|
x: &Tensor,
|
|
y: &Tensor,
|
|
train_idx: &[usize],
|
|
val_idx: &[usize],
|
|
) -> AutoMLResult<(Tensor, Tensor, Tensor, Tensor)> {
|
|
let x_train = index_tensor(x, train_idx)?;
|
|
let y_train = index_tensor(y, train_idx)?;
|
|
let x_val = index_tensor(x, val_idx)?;
|
|
let y_val = index_tensor(y, val_idx)?;
|
|
|
|
Ok((x_train, y_train, x_val, y_val))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_tensor::Device;
|
|
|
|
#[test]
|
|
fn test_find_min_basic() {
|
|
let device = Device::cpu();
|
|
let data = vec![3.0, 1.0, 4.0, 1.5, 9.0, 2.6];
|
|
let tensor = Tensor::from_slice(&data, &[6], &device).unwrap();
|
|
|
|
let min_val = find_min(&tensor).unwrap();
|
|
assert_eq!(min_val, 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_find_min_with_negatives() {
|
|
let device = Device::cpu();
|
|
let data = vec![-5.0, 3.0, -1.0, 4.0];
|
|
let tensor = Tensor::from_slice(&data, &[4], &device).unwrap();
|
|
|
|
let min_val = find_min(&tensor).unwrap();
|
|
assert_eq!(min_val, -5.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_min_returns_tensor() {
|
|
let device = Device::cpu();
|
|
let data = vec![3.0, 1.0, 4.0];
|
|
let tensor = Tensor::from_slice(&data, &[3], &device).unwrap();
|
|
|
|
let min_tensor = tensor_min(&tensor).unwrap();
|
|
let min_val = min_tensor.to_scalar::<f32>().unwrap();
|
|
assert_eq!(min_val, 1.0);
|
|
}
|
|
}
|